Skip to content

rLLM Advantage Estimator API (Experimental)

Reference for the rLLM-native advantage estimator interfaces, registry APIs, and related config types.

Registry and Collection

rllm.experimental.common.advantage.register_rllm_adv_estimator 

register_rllm_adv_estimator(name: str | rLLMAdvantageEstimator) -> Callable

Register a rLLM advantage estimator -- either built-in or custom.

Parameters:

Name Type Description Default
name str | rLLMAdvantageEstimator

Name of the advantage estimator.

 required
Source code in rllm/experimental/common/advantage.py 
25
26
27
28
29
30
31
32
33
34
35
36
def register_rllm_adv_estimator(name: str | rLLMAdvantageEstimator) -> Callable:
    """Register a rLLM advantage estimator -- either built-in or custom.

    Args:
        name: Name of the advantage estimator.
    """

    def decorator(func: Callable) -> Callable:
        RLLM_ADV_ESTIMATOR_REGISTRY[name] = func
        return func

    return decorator

rllm.experimental.common.advantage.get_rllm_adv_estimator 

get_rllm_adv_estimator(name: str | rLLMAdvantageEstimator) -> Callable

Get a rLLM advantage estimator by name.

Parameters:

Name Type Description Default
name str | rLLMAdvantageEstimator

Name of the advantage estimator.

 required
Source code in rllm/experimental/common/advantage.py 
39
40
41
42
43
44
45
46
47
def get_rllm_adv_estimator(name: str | rLLMAdvantageEstimator) -> Callable:
    """Get a rLLM advantage estimator by name.

    Args:
        name: Name of the advantage estimator.
    """
    if name not in RLLM_ADV_ESTIMATOR_REGISTRY:
        raise ValueError(f"Unknown advantage estimator {name}. If you have a custom advantage estimator, please register it using `register_rllm_adv_estimator`.")
    return RLLM_ADV_ESTIMATOR_REGISTRY[name]

rllm.experimental.common.advantage.collect_reward_and_advantage_from_trajectory_groups 

collect_reward_and_advantage_from_trajectory_groups(groups: list[TrajectoryGroup], algorithm_config: AlgorithmConfig, collect_advantage: bool = True) -> dict

Collect reward and advantage from trajectory groups. Return a dictionary of metrics. If collect_advantage is False, only collect rewards.

Parameters:

Name Type Description Default
groups list[TrajectoryGroup]

List of TrajectoryGroup objects

 required
algorithm_config AlgorithmConfig

Algorithm configuration

 required
collect_advantage bool

Whether to collect advantage

 True

Returns:

Type Description
dict

Dictionary of metrics
Source code in rllm/experimental/common/advantage.py 
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
def collect_reward_and_advantage_from_trajectory_groups(
    groups: list[TrajectoryGroup],
    algorithm_config: AlgorithmConfig,
    collect_advantage: bool = True,
) -> dict:
    """
    Collect reward and advantage from trajectory groups. Return a dictionary of metrics.
    If collect_advantage is False, only collect rewards.

    Args:
        groups: List of TrajectoryGroup objects
        algorithm_config: Algorithm configuration
        collect_advantage: Whether to collect advantage

    Returns:
        Dictionary of metrics
    """
    assert algorithm_config.stepwise_advantage_mode == "broadcast", "Only broadcast mode is supported in experimental unified trainer."

    advantages_by_role = defaultdict(list)
    rewards_by_role = defaultdict(list)
    traj_rewards_by_role = defaultdict(list)
    traj_groups_by_role = defaultdict(list)

    for group in groups:
        group_role = group.group_role
        has_precomputed_advantage = any(step.advantage is not None for traj in group.trajectories for step in traj.steps)

        if has_precomputed_advantage and algorithm_config.use_precomputed_advantage:
            # Precompute mode (e.g. OPD, SFT): always use pre-computed per-token advantages from the workflow.
            if collect_advantage:
                flattened_advantages = _collect_precomputed_advantages(group, group_role)
                advantages_by_role[group_role].extend(flattened_advantages)
        else:
            # RL mode: compute advantages from trajectory rewards.
            if collect_advantage and has_precomputed_advantage:
                logger.warning(f"[group={group_role}] Steps have pre-computed advantages but use_precomputed_advantage is False. Overwriting with {algorithm_config.estimator.value}.")

            assert all(traj.reward is not None for traj in group.trajectories), "Trajectory reward cannot be None in broadcast mode"
            traj_rewards = np.array([traj.reward for traj in group.trajectories])
            rewards_by_role[group_role].extend(traj_rewards)

            if collect_advantage:
                traj_groups_by_role[group_role].append(group)
                traj_rewards_by_role[group_role].append(traj_rewards)

    if collect_advantage:
        for group_role, traj_groups in traj_groups_by_role.items():
            advantage_fn = get_rllm_adv_estimator(algorithm_config.estimator_map.get(group_role, algorithm_config.estimator))
            traj_rewards = traj_rewards_by_role[group_role]
            adv_kwargs = _prepare_adv_estimator_input(traj_rewards, algorithm_config)
            advantages_by_group, _ = advantage_fn(**adv_kwargs)  # ignore the returns here
            assert len(advantages_by_group) == len(traj_groups), "length mismatch between advantages and trajectory groups"
            for traj_group, advantages_by_traj in zip(traj_groups, advantages_by_group, strict=True):
                assert len(advantages_by_traj) == len(traj_group.trajectories), "length mismatch between trajectory rewards and computed advantages"
                advantages_by_role[group_role].extend(np.asarray(advantages_by_traj).tolist())  # for metrics calculation
                for traj, advantage in zip(traj_group.trajectories, advantages_by_traj, strict=True):
                    for step in traj.steps:
                        step.advantage = advantage

    # reduce metrics by group
    final_metrics = {}
    for group_role, rewards in rewards_by_role.items():
        final_metrics[f"reward/{group_role}/mean"] = np.mean(rewards)
        final_metrics[f"reward/{group_role}/std"] = np.std(rewards)
        final_metrics[f"reward/{group_role}/max"] = np.max(rewards)
        final_metrics[f"reward/{group_role}/min"] = np.min(rewards)

    if collect_advantage:
        for group_role, advantages in advantages_by_role.items():
            final_metrics[f"advantage/{group_role}/mean"] = np.mean(advantages)
            final_metrics[f"advantage/{group_role}/std"] = np.std(advantages)
            final_metrics[f"advantage/{group_role}/max"] = np.max(advantages)
            final_metrics[f"advantage/{group_role}/min"] = np.min(advantages)
            final_metrics[f"advantage/{group_role}/fraction_zero"] = np.sum(np.abs(advantages) < 1e-8) / len(advantages)

    return final_metrics

Built-in Estimators

rllm.experimental.common.advantage.calculate_grpo_advantages 

calculate_grpo_advantages(rewards: list[ndarray], norm_adv_by_std_in_grpo=True, episilon=1e-06, **kwargs) -> tuple[list[np.ndarray], list[np.ndarray]]
Source code in rllm/experimental/common/advantage.py 
50
51
52
53
54
@register_rllm_adv_estimator(rLLMAdvantageEstimator.GRPO)
def calculate_grpo_advantages(rewards: list[np.ndarray], norm_adv_by_std_in_grpo=True, episilon=1e-6, **kwargs) -> tuple[list[np.ndarray], list[np.ndarray]]:
    advantages_by_group, returns_by_group = zip(*[calculate_grpo_advantages_per_group(group_rewards, norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo, episilon=episilon) for group_rewards in rewards], strict=True)

    return advantages_by_group, returns_by_group

rllm.experimental.common.advantage.calculate_reinforce_advantages 

calculate_reinforce_advantages(rewards: list[ndarray], **kwargs) -> tuple[list[np.ndarray], list[np.ndarray]]

REINFORCE: advantage = reward (no baseline)

Source code in rllm/experimental/common/advantage.py 
57
58
59
60
@register_rllm_adv_estimator(rLLMAdvantageEstimator.REINFORCE)
def calculate_reinforce_advantages(rewards: list[np.ndarray], **kwargs) -> tuple[list[np.ndarray], list[np.ndarray]]:
    """REINFORCE: advantage = reward (no baseline)"""
    return rewards, rewards

rllm.experimental.common.advantage.calculate_reinforce_plus_plus_baseline_advantages 

calculate_reinforce_plus_plus_baseline_advantages(rewards: list[ndarray], epsilon=1e-06, **kwargs) -> tuple[list[np.ndarray], list[np.ndarray]]

REINFORCE++ baseline estimator.

In line with Verl's REINFORCE++ baseline logic for grouped rollouts: 1. Use per-group mean baseline when group size > 1, else baseline = 0. 2. Whiten centered scores using role-level batch statistics.

Source code in rllm/experimental/common/advantage.py 
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
@register_rllm_adv_estimator(rLLMAdvantageEstimator.REINFORCE_PLUS_PLUS_BASELINE)
def calculate_reinforce_plus_plus_baseline_advantages(rewards: list[np.ndarray], epsilon=1e-6, **kwargs) -> tuple[list[np.ndarray], list[np.ndarray]]:
    """REINFORCE++ baseline estimator.

    In line with Verl's REINFORCE++ baseline logic for grouped rollouts:
    1. Use per-group mean baseline when group size > 1, else baseline = 0.
    2. Whiten centered scores using role-level batch statistics.
    """
    if len(rewards) == 0:
        return [], []

    centered_rewards_by_group: list[np.ndarray] = []
    for group_rewards in rewards:
        centered_rewards_by_group.append(group_rewards - np.mean(group_rewards))

    all_centered_rewards = np.concatenate(centered_rewards_by_group)
    batch_std = np.std(all_centered_rewards)

    advantages_by_group = [centered_rewards / (batch_std + epsilon) for centered_rewards in centered_rewards_by_group]

    return advantages_by_group, advantages_by_group

rllm.experimental.common.advantage.calculate_rloo_advantages 

calculate_rloo_advantages(rewards: list[ndarray], **kwargs) -> tuple[list[np.ndarray], list[np.ndarray]]

Reinforce Leave-one-out (RLOO): https://arxiv.org/abs/2402.14740

Source code in rllm/experimental/common/advantage.py 
86
87
88
89
90
@register_rllm_adv_estimator(rLLMAdvantageEstimator.RLOO)
def calculate_rloo_advantages(rewards: list[np.ndarray], **kwargs) -> tuple[list[np.ndarray], list[np.ndarray]]:
    """Reinforce Leave-one-out (RLOO): https://arxiv.org/abs/2402.14740"""
    advantages_by_group, returns_by_group = zip(*[calculate_rloo_advantages_per_group(group_rewards) for group_rewards in rewards], strict=True)
    return advantages_by_group, returns_by_group

Supporting Config Types

rllm.experimental.common.config.rLLMAdvantageEstimator

Bases: str, Enum

A unified advantage estimator for rLLM. Work with both tinker and verl backends at the expense of losing some flexibility. TODO(listar2000): add more estimators.

Source code in rllm/experimental/common/config.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
class rLLMAdvantageEstimator(str, Enum):
    """
    A unified advantage estimator for rLLM. Work with both `tinker` and `verl` backends at the expense of
    losing some flexibility.
    TODO(listar2000): add more estimators.
    """

    GRPO = "grpo"
    REINFORCE = "reinforce"
    REINFORCE_PLUS_PLUS_BASELINE = "reinforce_plus_plus_baseline"
    RLOO = "rloo"
    OTHER = "other"

    @classmethod
    def _missing_(cls, value: object) -> "rLLMAdvantageEstimator":
        return cls.OTHER

rllm.experimental.common.config.AlgorithmConfig dataclass

Configuration for algorithm parameters.

Source code in rllm/experimental/common/config.py 
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
@dataclass
class AlgorithmConfig:
    """Configuration for algorithm parameters."""

    use_rllm: bool = False  # This is ignored (assumed True) for tinker backend.
    estimator: rLLMAdvantageEstimator = rLLMAdvantageEstimator.GRPO
    estimator_map: dict[str, rLLMAdvantageEstimator | str] = field(default_factory=dict)
    # TODO(listar2000): eventually we will remove the `per_step` mode all-together. Now we keep it for backward compatibility.
    stepwise_advantage_mode: Literal["broadcast", "per_step"] = "broadcast"
    norm_adv_by_std_in_grpo: bool = True
    # When True, always use pre-computed step.advantage from the workflow and skip
    # advantage computation (GRPO/REINFORCE). Steps missing advantages default to 0.0.
    # When False (default), always compute advantages normally.
    use_precomputed_advantage: bool = False
    # for tinker backend only
    loss_fn: Literal["importance_sampling", "ppo", "cispo", "dro", "cross_entropy"] | None = None
    lr_schedule: Literal["linear", "cosine", "constant"] = "constant"
    warmup_steps_ratio: float = 0.0

    @classmethod
    def from_config(cls, config: DictConfig) -> "AlgorithmConfig":
        """Create an AlgorithmConfig from a dictionary configuration.

        Args:
            config: Dictionary configuration.
        Returns:
            AlgorithmConfig: The AlgorithmConfig built from the configuration.
        """
        return cls(
            estimator=rLLMAdvantageEstimator(config.algorithm.adv_estimator),
            stepwise_advantage_mode=config.rllm.stepwise_advantage.mode,
            norm_adv_by_std_in_grpo=config.rllm.stepwise_advantage.get("norm_adv_by_std_in_grpo", True),
            use_rllm=config.rllm.stepwise_advantage.get("use_rllm", False),
            use_precomputed_advantage=config.rllm.algorithm.get("use_precomputed_advantage", False),
            loss_fn=config.rllm.algorithm.get("loss_fn", None),
            lr_schedule=config.rllm.algorithm.get("lr_schedule", "constant"),
            warmup_steps_ratio=config.rllm.algorithm.get("warmup_steps_ratio", 0.0),
        )

    def __post_init__(self):
        if self.stepwise_advantage_mode == "per_step":
            from warnings import warn

            warn(
                "The `per_step` mode is deprecated in experimental unified trainer. Set to `broadcast` mode automatically. Please either use the legacy trainers (`agent_workflow_trainer` for `Verl` or `tinker_workflow_trainer` for `Tinker`) with the `per_step` configuration. Or manually pass in a hook with the implementation of `per_step` advantage computation logic.",
                DeprecationWarning,
                stacklevel=2,
            )
            self.stepwise_advantage_mode = "broadcast"

from_config classmethod 

from_config(config: DictConfig) -> AlgorithmConfig

Create an AlgorithmConfig from a dictionary configuration.

Parameters:

Name Type Description Default
config DictConfig

Dictionary configuration.

 required

Returns: AlgorithmConfig: The AlgorithmConfig built from the configuration.

Source code in rllm/experimental/common/config.py 
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
@classmethod
def from_config(cls, config: DictConfig) -> "AlgorithmConfig":
    """Create an AlgorithmConfig from a dictionary configuration.

    Args:
        config: Dictionary configuration.
    Returns:
        AlgorithmConfig: The AlgorithmConfig built from the configuration.
    """
    return cls(
        estimator=rLLMAdvantageEstimator(config.algorithm.adv_estimator),
        stepwise_advantage_mode=config.rllm.stepwise_advantage.mode,
        norm_adv_by_std_in_grpo=config.rllm.stepwise_advantage.get("norm_adv_by_std_in_grpo", True),
        use_rllm=config.rllm.stepwise_advantage.get("use_rllm", False),
        use_precomputed_advantage=config.rllm.algorithm.get("use_precomputed_advantage", False),
        loss_fn=config.rllm.algorithm.get("loss_fn", None),
        lr_schedule=config.rllm.algorithm.get("lr_schedule", "constant"),
        warmup_steps_ratio=config.rllm.algorithm.get("warmup_steps_ratio", 0.0),
    )

Per-Group Algo Helpers

rllm.experimental.common.rl_algo.calculate_grpo_advantages_per_group 

calculate_grpo_advantages_per_group(rewards: ndarray, norm_adv_by_std_in_grpo=True, episilon=1e-06) -> tuple[np.ndarray, np.ndarray]
Source code in rllm/experimental/common/rl_algo.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
def calculate_grpo_advantages_per_group(rewards: np.ndarray, norm_adv_by_std_in_grpo=True, episilon=1e-6) -> tuple[np.ndarray, np.ndarray]:
    if len(rewards) <= 1:
        group_mean, group_std = 0.0, 1.0
    else:
        group_mean = np.mean(rewards)
        group_std = np.std(rewards)

    if norm_adv_by_std_in_grpo:
        advantages = (rewards - group_mean) / (group_std + episilon)
    else:
        advantages = rewards - group_mean

    return advantages, advantages

rllm.experimental.common.rl_algo.calculate_rloo_advantages_per_group 

calculate_rloo_advantages_per_group(rewards: ndarray) -> tuple[np.ndarray, np.ndarray]
Source code in rllm/experimental/common/rl_algo.py 
21
22
23
24
25
26
27
def calculate_rloo_advantages_per_group(rewards: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    num_trajs = len(rewards)
    if num_trajs <= 1:
        return rewards

    advantages = num_trajs / (num_trajs - 1) * (rewards - rewards.mean())
    return advantages, advantages

UnifiedTrainer Entry Point

traj_group_adv_estimator_map is provided through trainer construction kwargs and wired into AlgorithmConfig.estimator_map in UnifiedTrainer.

rllm.experimental.unified_trainer.UnifiedTrainer

Unified trainer for backend-agnostic training.

This trainer uses an async-prioritized design where the core pipeline methods are async. This accommodates backends that naturally use async operations (like Tinker) while still supporting sync backends.

The main fit() method remains sync for ease of use, but internally runs the async training loop in a dedicated event loop thread.

Source code in rllm/experimental/unified_trainer.py 
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
class UnifiedTrainer:
    """Unified trainer for backend-agnostic training.

    This trainer uses an async-prioritized design where the core pipeline methods
    are async. This accommodates backends that naturally use async operations
    (like Tinker) while still supporting sync backends.

    The main `fit()` method remains sync for ease of use, but internally runs
    the async training loop in a dedicated event loop thread.
    """

    def __init__(
        self,
        backend_cls: type[BackendProtocol],
        config: DictConfig,
        workflow_class: type[Workflow],
        train_dataset: Dataset | None = None,
        val_dataset: Dataset | None = None,
        workflow_args: dict | None = None,
        backend_args: dict | None = None,
        *,
        traj_grouping_hook: Callable | None = None,
        traj_group_adv_estimator_map: dict | None = None,
        **kwargs,
    ):
        """Initialize the UnifiedTrainer."""
        self.workflow_class = workflow_class
        self.workflow_args = workflow_args or {}
        self.train_dataset = train_dataset
        self.val_dataset = val_dataset

        # initializing and validating common configs
        self.config = config
        self.rllm_config = config.rllm

        # Read user-defined hooks from kwargs
        self.traj_grouping_hook = traj_grouping_hook or _default_traj_grouping_hook
        # Extract the TrajectoryGroup-specific estimator from kwargs
        self.traj_group_adv_estimator_map = traj_group_adv_estimator_map or {}

        self.backend = backend_cls(config=config, **(backend_args or {}))

        self._validate_and_setup_configs()
        self._setup_logging()

        rollout_engine: RolloutEngine = self.backend.init_rollout_engine(
            cf_config=self.cf_config,
            transform_config=self.transform_config,
            rs_config=self.rs_config,
            algorithm_config=self.algorithm_config,
        )
        self.agent_workflow_engine = UnifiedWorkflowEngine(
            workflow_cls=self.workflow_class,
            workflow_args=self.workflow_args,
            rollout_engine=rollout_engine,
            config=self.config,
            n_parallel_tasks=self.rllm_config.workflow.n_parallel_tasks,
            retry_limit=self.rllm_config.workflow.retry_limit,
            raise_on_error=self.rllm_config.workflow.raise_on_error,
            episode_logger=self.episode_logger,
        )

        self.tokenizer = None
        if hasattr(self.backend, "tokenizer"):
            self.tokenizer = self.backend.tokenizer

    def _validate_and_setup_configs(self):
        """Validate and setup common configs."""
        # validate common, backend-agnostic configs
        assert self.rllm_config is not None, "rLLM config is not set"
        # if the traj_group_adv_estimator_map is given, the user must turn `use_rllm` to True
        if self.traj_group_adv_estimator_map and not self.rllm_config.algorithm.get("use_rllm", False):
            raise ValueError("If `traj_group_adv_estimator_map` is given, the user must explicitly turn `rllm.algorithm.use_rllm` to True")

        if self.rllm_config.rejection_sample.multiplier != 1:
            assert self.rllm_config.rejection_sample.enable is True, "rejection sampling is disabled, but rejection_sample.multiplier is not 1"

        # validate backend-specific configs
        self.backend.validate_config()

        # compact filtering config (used for filtering out episodes that are not valid)
        self.cf_config = CompactFilteringConfig.from_config(self.rllm_config.compact_filtering)

        # transform config (used for transforming episodes to trajectory groups)
        self.transform_config = TransformConfig(broadcast=self.rllm_config.stepwise_advantage.mode == "broadcast")

        # rejection sampling config (used for rejection sampling)
        rs_mode = "episode" if self.rllm_config.rejection_sample.enable else "none"

        self.rs_config = RejectionSamplingConfig(
            mode=rs_mode,
            min_partial_solve_tasks=self.rllm_config.rejection_sample.min_partial_solve_tasks,
            min_trajs_per_group=self.rllm_config.rejection_sample.min_trajs_per_group,
        )

        # algorithm config (used for rLLM-native advantage computation)
        self.algorithm_config = AlgorithmConfig(
            estimator=self.rllm_config.algorithm.adv_estimator,
            estimator_map=self.traj_group_adv_estimator_map,  # TODO(listar2000): see if we can make this configurable in config as well
            stepwise_advantage_mode=self.rllm_config.stepwise_advantage.mode,
            norm_adv_by_std_in_grpo=self.rllm_config.stepwise_advantage.get("norm_adv_by_std_in_grpo", True),
            use_rllm=self.rllm_config.algorithm.get("use_rllm", False),
            use_precomputed_advantage=self.rllm_config.algorithm.get("use_precomputed_advantage", False),
            loss_fn=self.rllm_config.algorithm.get("loss_fn", None),
            lr_schedule=self.rllm_config.algorithm.get("lr_schedule", "constant"),
            warmup_steps_ratio=self.rllm_config.algorithm.get("warmup_steps_ratio", 0.0),
        )

    def _setup_logging(self):
        """Setup up both the tracking and episode logging."""
        # create episode logger if enabled in config
        self.episode_logger = None
        if self.rllm_config.episode_logging.get("log_episodes", False):
            episode_log_dir = self.rllm_config.episode_logging.get(
                "episode_log_dir",
                f"logs/{self.rllm_config.trainer.project_name}/{self.rllm_config.trainer.experiment_name}",
            )
            self.episode_logger = EpisodeLogger(base_dir=episode_log_dir, subdirectory="episodes")

        source_metadata = extract_source_metadata(
            workflow_class=self.workflow_class,
            workflow_args=self.workflow_args,
        )

        self.logger = Tracking(
            project_name=self.rllm_config.trainer.project_name,
            experiment_name=self.rllm_config.trainer.experiment_name,
            default_backend=self.rllm_config.trainer.logger,
            config=OmegaConf.to_container(self.config, resolve=True),
            source_metadata=source_metadata,
        )

    # =========================================================================
    # Main training loop methods
    # =========================================================================

    def fit(self):
        """Main training loop (sync entry point)."""
        asyncio.run(self.fit_async())

    async def fit_async(self) -> None:
        """Public async entry point for the full training process."""
        # initialize the UnifiedWorkflowEngine (init the workflow pool)
        await self.agent_workflow_engine.initialize_pool()

        trainer_state = TrainerState()

        await self.backend.on_train_start(trainer_state)

        if self.rllm_config.trainer.get("val_before_train", True):
            val_metrics = await self._validate_async(trainer_state)
            pprint(f"Initial validation metrics: {val_metrics}")
            if self.rllm_config.trainer.get("val_only", False):
                return

        # we start from step (1 + original start batch index)
        trainer_state.global_step += 1

        # Run the training loop
        await self._fit_async(trainer_state)

        await self.backend.on_train_end(trainer_state)

    async def _fit_async(self, trainer_state: TrainerState) -> None:
        """Internal async main training loop."""
        train_dataloader: Iterable = self.backend.get_dataloader(self.train_dataset, trainer_state)
        break_via_total_batches = False  # used to break the training loop via the `total_batches` parameter
        use_total_batches = self.rllm_config.trainer.get("total_batches") is not None and self.rllm_config.trainer.total_batches > 0

        if use_total_batches:
            trainer_state.total_steps = self.rllm_config.trainer.total_batches
        else:
            trainer_state.total_steps = len(train_dataloader) * self.rllm_config.trainer.total_epochs

        for epoch in range(self.rllm_config.trainer.total_epochs):
            # recursively break through the outer loop
            if break_via_total_batches:
                break

            pprint(f"epoch {epoch}, step {trainer_state.global_step} started")
            trainer_state.epoch = epoch
            await self.backend.on_epoch_start(trainer_state)

            for batch in train_dataloader:
                trainer_state.reset_batch()

                await self.backend.on_batch_start(trainer_state)
                with simple_timer("total_step", trainer_state.timing_dict):
                    await self._train_batch_async(batch, trainer_state)
                await self.backend.on_batch_end(trainer_state)

                self.logger.log(
                    data=trainer_state.metrics,
                    step=trainer_state.global_step,
                    episodes=trainer_state.episodes,
                    trajectory_groups=trainer_state.trajectory_groups,
                )

                # if the config specifies the `total_batches` parameter, then we check if we should stop
                if use_total_batches and trainer_state.global_step >= self.rllm_config.trainer.total_batches:
                    break_via_total_batches = True
                    break

                # periodic validation
                if self.rllm_config.trainer.test_freq > 0 and trainer_state.global_step % self.rllm_config.trainer.test_freq == 0:
                    await self._validate_async(trainer_state)

                trainer_state.global_step += 1

            await self.backend.on_epoch_end(trainer_state)

        # final validation after training
        if self.rllm_config.trainer.test_freq > 0:
            val_metrics = await self._validate_async(trainer_state)
            pprint(f"Final validation metrics: {val_metrics}")

    async def _train_batch_async(self, batch: Any, trainer_state: TrainerState) -> None:
        """Train a batch (async implementation)."""
        self.agent_workflow_engine.set_training_step(trainer_state.global_step, mode="train", epoch=trainer_state.epoch)

        # stage 1: generate episodes (async) and collect metrics (sync)
        trainer_state.episodes = await self.backend.generate_episodes(batch, agent_workflow_engine=self.agent_workflow_engine, is_validation=False)
        if not trainer_state.has_episodes:
            return

        workflow_metrics, termination_counts = self._collect_workflow_metrics_from_episodes(trainer_state.episodes)

        # stage 2: transform episodes to trajectory groups (sync)
        trajectory_groups, transform_metrics = transform_episodes_to_trajectory_groups(trainer_state.episodes, self.transform_config, self.cf_config, traj_grouping_hook=self.traj_grouping_hook)
        trainer_state.trajectory_groups = trajectory_groups
        trainer_state.metrics.update(transform_metrics)

        # stage 3: apply rejection sampling (sync)
        filtered_groups, filtered_episodes, rs_metrics = apply_rejection_sampling_and_filtering(
            trainer_state.episodes,
            trainer_state.trajectory_groups,
            self.rs_config,
            trainer_state.rs_state,
        )
        trainer_state.metrics.update(rs_metrics)
        trainer_state.trajectory_groups = filtered_groups
        trainer_state.episodes = filtered_episodes
        if not trainer_state.has_trajectory_groups:
            return

        # stage 4: transform rllm-native data structures to backend-specific format (sync)
        backend_batch = self.backend.transform_to_backend_batch(trainer_state)
        trainer_state.backend_batch = backend_batch

        # stage 5: process backend batch (async) - compute log probs, critic values, etc.
        await self.backend.process_backend_batch(trainer_state)
        assert trainer_state.has_backend_batch, "Backend batch is not transformed or processed successfully"

        # stage 6: compute advantages (async)
        await self.backend.compute_advantages(trainer_state, self.algorithm_config)

        # stage 7: update policy (async)
        await self.backend.update_policy(trainer_state)

        # stage 8: cleanup, logging, visualization, etc. (sync)
        if self.tokenizer is not None:
            visualize_trajectory_last_steps(
                trainer_state.trajectory_groups,
                tokenizer=self.tokenizer,
                max_steps_to_visualize=2,
                show_workflow_metadata=True,
            )

        for key, value in workflow_metrics.items():
            trainer_state.metrics[f"batch/{key}"] = np.mean(value)

        total_counts = max(sum(termination_counts.values()), 1)
        for r in TerminationReason:
            trainer_state.metrics[f"batch/termination_reason/{r.value}"] = termination_counts[r.value] / total_counts

    async def _validate_async(self, trainer_state: TrainerState) -> dict:
        """Validate the model (async implementation)."""
        n_val_samples = self.rllm_config.rollout.n_val
        val_metrics = defaultdict(list)

        if not await self.backend.on_validation_start(trainer_state):
            return {}
        # manually manage the testing time
        test_begin = time.perf_counter()
        self.agent_workflow_engine.set_training_step(trainer_state.global_step, mode="val", epoch=trainer_state.epoch)

        is_correct_lst, uid_lst, data_source_lst = [], [], []
        workflow_metrics_by_source = defaultdict(lambda: defaultdict(list))

        val_dataloader: Iterable = self.backend.get_dataloader(self.val_dataset, trainer_state)
        for batch in val_dataloader:
            # Generate episodes and transform to trajectory groups
            val_episodes = await self.backend.generate_episodes(batch, agent_workflow_engine=self.agent_workflow_engine, is_validation=True)
            val_trajectory_groups, transform_metrics = transform_episodes_to_trajectory_groups(val_episodes, self.transform_config, self.cf_config, traj_grouping_hook=self.traj_grouping_hook)
            reward_metrics = collect_reward_and_advantage_from_trajectory_groups(val_trajectory_groups, self.algorithm_config, collect_advantage=False)

            is_correct_lst.extend([episode.is_correct for episode in val_episodes])
            uid_lst.extend([episode.task_id for episode in val_episodes])

            data_sources = [episode.info.get("data_source", "unknown") for episode in val_episodes]
            data_source_lst.extend(data_sources)

            for episode, data_source in zip(val_episodes, data_sources, strict=True):
                for key, value in episode.metrics.items():
                    workflow_metrics_by_source[data_source][key].append(float(value))

            for key, value in (transform_metrics | reward_metrics).items():
                val_metrics[f"val/{key}"].append(value)

        test_end = time.perf_counter()
        val_metrics["time/testing"] = test_end - test_begin
        is_correct_array = np.array(is_correct_lst)
        uid_array = np.array(uid_lst)
        data_source_array = np.array(data_source_lst)

        for data_source in np.unique(data_source_array):
            pass_rates = defaultdict(list)

            data_source_mask = data_source_array == data_source
            is_correct_data_source = is_correct_array[data_source_mask]
            uids_data_source = uid_array[data_source_mask]

            for is_correct, uid in zip(is_correct_data_source, uids_data_source, strict=False):
                pass_rates[uid].append(is_correct)

            val_metrics[f"val/{data_source}/pass@1"] = np.mean(is_correct_data_source)
            val_metrics[f"val/{data_source}/pass@{n_val_samples}"] = np.mean([1 if any(pass_rate) else 0 for pass_rate in pass_rates.values()])

            # Add workflow metrics for this data source
            if data_source in workflow_metrics_by_source:
                for key, values in workflow_metrics_by_source[data_source].items():
                    if values:
                        val_metrics[f"val/{data_source}/{key}"] = np.mean(values)

        # post-process the val metrics to reduce any "list values" into scalars
        reduce_metrics_lists(val_metrics)
        self.logger.log(data=val_metrics, step=trainer_state.global_step)
        await self.backend.on_validation_end(trainer_state)
        return val_metrics

    def shutdown(self):
        """Shutdown the trainer and cleanup resources."""
        if hasattr(self, "agent_workflow_engine") and self.agent_workflow_engine is not None:
            self.agent_workflow_engine.shutdown()
        self.backend.shutdown()

        # Explicitly finish the logger to prevent hang in __del__ during garbage collection
        if hasattr(self, "logger") and self.logger is not None:
            self.logger.finish()

    # =========================================================================
    # Helper functions
    # =========================================================================
    def _collect_workflow_metrics_from_episodes(self, episodes: list[Episode]) -> tuple[dict, Counter]:
        workflow_metrics = defaultdict(list)
        termination_counts = Counter()
        for episode in episodes:
            for k, v in episode.metrics.items():
                workflow_metrics[k].append(v)
            if episode.termination_reason is not None:
                termination_counts[episode.termination_reason.value] += 1
        # reduce the metrics to a scalar value, with error handling
        reduced_workflow_metrics = {}
        for k, v in workflow_metrics.items():
            try:
                reduced_workflow_metrics[k] = np.mean(v)
            except Exception:
                continue
        return reduced_workflow_metrics, termination_counts

__init__ 

__init__(backend_cls: type[BackendProtocol], config: DictConfig, workflow_class: type[Workflow], train_dataset: Dataset | None = None, val_dataset: Dataset | None = None, workflow_args: dict | None = None, backend_args: dict | None = None, *, traj_grouping_hook: Callable | None = None, traj_group_adv_estimator_map: dict | None = None, **kwargs)

Initialize the UnifiedTrainer.

Source code in rllm/experimental/unified_trainer.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
def __init__(
    self,
    backend_cls: type[BackendProtocol],
    config: DictConfig,
    workflow_class: type[Workflow],
    train_dataset: Dataset | None = None,
    val_dataset: Dataset | None = None,
    workflow_args: dict | None = None,
    backend_args: dict | None = None,
    *,
    traj_grouping_hook: Callable | None = None,
    traj_group_adv_estimator_map: dict | None = None,
    **kwargs,
):
    """Initialize the UnifiedTrainer."""
    self.workflow_class = workflow_class
    self.workflow_args = workflow_args or {}
    self.train_dataset = train_dataset
    self.val_dataset = val_dataset

    # initializing and validating common configs
    self.config = config
    self.rllm_config = config.rllm

    # Read user-defined hooks from kwargs
    self.traj_grouping_hook = traj_grouping_hook or _default_traj_grouping_hook
    # Extract the TrajectoryGroup-specific estimator from kwargs
    self.traj_group_adv_estimator_map = traj_group_adv_estimator_map or {}

    self.backend = backend_cls(config=config, **(backend_args or {}))

    self._validate_and_setup_configs()
    self._setup_logging()

    rollout_engine: RolloutEngine = self.backend.init_rollout_engine(
        cf_config=self.cf_config,
        transform_config=self.transform_config,
        rs_config=self.rs_config,
        algorithm_config=self.algorithm_config,
    )
    self.agent_workflow_engine = UnifiedWorkflowEngine(
        workflow_cls=self.workflow_class,
        workflow_args=self.workflow_args,
        rollout_engine=rollout_engine,
        config=self.config,
        n_parallel_tasks=self.rllm_config.workflow.n_parallel_tasks,
        retry_limit=self.rllm_config.workflow.retry_limit,
        raise_on_error=self.rllm_config.workflow.raise_on_error,
        episode_logger=self.episode_logger,
    )

    self.tokenizer = None
    if hasattr(self.backend, "tokenizer"):
        self.tokenizer = self.backend.tokenizer

fit 

fit()

Main training loop (sync entry point).

Source code in rllm/experimental/unified_trainer.py 
217
218
219
def fit(self):
    """Main training loop (sync entry point)."""
    asyncio.run(self.fit_async())

fit_async async 

fit_async() -> None

Public async entry point for the full training process.

Source code in rllm/experimental/unified_trainer.py 
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
async def fit_async(self) -> None:
    """Public async entry point for the full training process."""
    # initialize the UnifiedWorkflowEngine (init the workflow pool)
    await self.agent_workflow_engine.initialize_pool()

    trainer_state = TrainerState()

    await self.backend.on_train_start(trainer_state)

    if self.rllm_config.trainer.get("val_before_train", True):
        val_metrics = await self._validate_async(trainer_state)
        pprint(f"Initial validation metrics: {val_metrics}")
        if self.rllm_config.trainer.get("val_only", False):
            return

    # we start from step (1 + original start batch index)
    trainer_state.global_step += 1

    # Run the training loop
    await self._fit_async(trainer_state)

    await self.backend.on_train_end(trainer_state)

shutdown 

shutdown()

Shutdown the trainer and cleanup resources.

Source code in rllm/experimental/unified_trainer.py 
421
422
423
424
425
426
427
428
429
def shutdown(self):
    """Shutdown the trainer and cleanup resources."""
    if hasattr(self, "agent_workflow_engine") and self.agent_workflow_engine is not None:
        self.agent_workflow_engine.shutdown()
    self.backend.shutdown()

    # Explicitly finish the logger to prevent hang in __del__ during garbage collection
    if hasattr(self, "logger") and self.logger is not None:
        self.logger.finish()