fix: account for multi-turn chunks in agentic LR scheduler budget (#407)

[dashitongzhi]

Problem

When using AgentNativeStepEnvManager for step-level agentic training, the LR scheduler exhausts its step budget far before all pipeline steps complete, causing the learning rate to drop to zero mid-training.

Root cause: PPOConfig.set_max_steps() computes total optimizer steps based on rollout_batch_size (number of trajectories), but AgentNativeStepEnvManager.formulate_rollouts() creates one training sample per turn — so the actual number of optimizer steps per pipeline step is much higher than budgeted.

Example from the field: With 4 trajectories × ~10 turns each = ~40 training samples per pipeline step. With backward_batch_size=4, that's ~10 optimizer steps per pipeline step — not the 1 that the scheduler was budgeted for. In a 200-step run, LR hit zero at step 123 (38.5% of training with zero LR).

Fixes #407.

Solution

This PR overrides set_max_steps in AgenticConfig to account for multi-turn chunking:

1. New config field: estimated_chunks_per_traj (default 0 = auto-detect)
2. Auto-detection: When not explicitly set, estimates from custom_envs.max_steps using max_steps // 2 as a conservative midpoint
3. Explicit override: Users can set estimated_chunks_per_traj for precise control
4. Backward compatible: When max_steps <= 1 in env configs (single-turn), falls back to multiplier of 1 — identical to parent behavior

How it works

# In AgenticConfig.set_max_steps():
chunks_multiplier = self._get_chunks_per_traj_estimate()

self.actor_train.training_args.max_steps = max(1, (
    max_steps * rollout_batch_size * num_return_sequences
    * ppo_epochs * chunks_multiplier // actor_backward_batch_size
))

User-facing config

# Auto-detect (default) — reads max_steps from custom_envs
estimated_chunks_per_traj: 0

# Explicit override for precise control
estimated_chunks_per_traj: 10

Files Changed

• roll/pipeline/agentic/agentic_config.py — Added estimated_chunks_per_traj field, _get_chunks_per_traj_estimate() method, and set_max_steps() override

Testing

The fix is backward-compatible: for single-turn environments (max_steps=1), the auto-detect returns 1 and behavior is identical to the parent class.

[Copilot]

Pull request overview

This PR fixes premature LR-scheduler exhaustion during step-level agentic training by making AgenticConfig.set_max_steps() account for the fact that some env managers generate one training sample per turn (chunk) rather than per trajectory.

Changes:

• Added estimated_chunks_per_traj to let users explicitly scale the optimizer-step budget for multi-turn (chunked) rollouts.
• Implemented _get_chunks_per_traj_estimate() to auto-estimate the chunk multiplier from custom_envs.
• Overrode set_max_steps() in AgenticConfig to multiply actor/critic training_args.max_steps by the chunk multiplier.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[CLAassistant]

All committers have signed the CLA.