[none0663]

What does this PR do?

fix bug for remax when the rollout mode is async, as metioned in #2551

[gemini-code-assist]

Code Review

This pull request correctly fixes a bug in the ReMax advantage estimation for asynchronous rollout mode by ensuring the correct sequence generator is used. My review identified one high-severity issue related to code duplication, which could pose a maintenance risk. I've recommended refactoring the duplicated logic into a helper method to improve maintainability and prevent future bugs.

[gemini-code-assist]

This logic for selecting the sequence generator based on self.async_rollout_mode is duplicated from lines 1177-1180. This duplication creates a maintenance burden and increases the risk of introducing bugs if the logic needs to be updated in the future and one of the instances is missed.

To mitigate this risk and improve code clarity, I strongly recommend refactoring this repeated logic into a single helper method.

For example:

# In RayPPOTrainer class
def _get_sequence_generator(self):
    """Returns the appropriate sequence generator based on the rollout mode."""
    if self.async_rollout_mode:
        return self.async_rollout_manager
    return self.actor_rollout_wg

# Then the call sites can be simplified:
generator = self._get_sequence_generator()
gen_batch_output = generator.generate_sequences(gen_batch)
# ... and for the baseline:
gen_baseline_output = generator.generate_sequences(gen_baseline_batch)

Applying this pattern would make the code more robust and easier to maintain.