[dubin555]

Problem

The OTB (Optimal Token Baseline) and TIR_OTB dispatch paths in ray_trainer.py do not pass old_log_probs to the estimator functions, causing a TypeError at runtime. Additionally, the two OTB estimator functions in core_algos.py lack **kwargs to absorb extra keys in adv_kwargs (like config and reward_baselines), causing a second TypeError. Together these make the entire OTB feature unusable.

Fix

1. Pass the old_log_probs parameter to the OTB/TIR_OTB estimator calls in ray_trainer.py.
2. Add **kwargs to compute_optimal_token_baseline_advantage and compute_multi_turn_optimal_token_baseline_advantage in core_algos.py, consistent with all other advantage estimators.

Testing

• CPU validation confirms both function signatures now match
• The OTB code path is now callable without TypeError

[gemini-code-assist]

Code Review

This pull request correctly addresses a TypeError in the Optimal Token Baseline (OTB) feature by passing the missing old_log_probs argument. However, my review has identified another underlying TypeError that will still prevent this feature from being usable. I have added a critical review comment detailing this issue and a recommended fix.

[gemini-code-assist]

While adding old_log_probs is the correct fix for the TypeError mentioned in the PR description, the subsequent call to adv_estimator_fn(**adv_kwargs) on line 221 will still raise a TypeError.

The adv_kwargs dictionary contains keys like config (and potentially reward_baselines) which are not accepted by the target functions compute_optimal_token_baseline_advantage and compute_multi_turn_optimal_token_baseline_advantage in verl/trainer/ppo/core_algos.py. These functions are missing a **kwargs parameter to accept unused arguments, unlike other advantage estimators in the same file.

To fully fix the OTB feature, you should modify the signatures of those two functions in verl/trainer/ppo/core_algos.py to accept arbitrary keyword arguments. For example:

def compute_optimal_token_baseline_advantage(
    token_level_rewards: torch.Tensor,
    response_mask: torch.Tensor,
    index: np.ndarray,
    old_log_probs: torch.Tensor,
    sum_pi_squared: torch.Tensor,
    rollout_is_weights: torch.Tensor = None,
    handle_zero_tail: bool = True,
    epsilon: float = 1e-8,
    **kwargs,  # Add this
) -> tuple[torch.Tensor, torch.Tensor]:

A similar change is needed for compute_multi_turn_optimal_token_baseline_advantage. This will prevent the new TypeError and make the OTB feature fully functional.

[dubin555]

Good catch — you're right that adv_kwargs contains extra keys like config and reward_baselines that would cause a second TypeError. This was addressed in the follow-up commit (2df3707): both compute_optimal_token_baseline_advantage and compute_multi_turn_optimal_token_baseline_advantage now accept **kwargs to absorb unused arguments, consistent with the other advantage estimators in core_algos.py.

I've also updated the PR description to reflect both fixes.

[ETOgaosion]

I think it makes sense, for future extension, we can add *kwargs and this seems to be a bug.