Skip to content

[algo, trainer] fix: pass missing old_log_probs to OTB estimators #5615

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ETOgaosion merged 2 commits into from
Mar 18, 2026
Merged

[algo, trainer] fix: pass missing old_log_probs to OTB estimators #5615

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5530999
fix(otb): pass old_log_probs to optimal token baseline estimators
dubin555 Mar 16, 2026
2df3707
fix(otb): add **kwargs to OTB estimator functions for adv_kwargs comp…
dubin555 Mar 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions verl/trainer/ppo/ray_trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -211,6 +211,8 @@ def compute_advantage(
"Please set actor.calculate_sum_pi_squared=True in config."
)
adv_kwargs["sum_pi_squared"] = data.batch["sum_pi_squared"]
# old_log_probs needed for path-variance proxy: w_t = 1 - 2*exp(old_log_probs) + sum_pi_squared
adv_kwargs["old_log_probs"] = data.batch["old_log_probs"]
# Get pre-computed rollout IS weights if available
rollout_is_weights = data.batch.get("rollout_is_weights", None)
adv_kwargs["rollout_is_weights"] = rollout_is_weights
Comment on lines 213 to 218

@gemini-code-assist gemini-code-assist Bot Mar 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

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.

Expand Down