[Refactor] Move train data preparation into AgentLoop and TrainingController

[YanhuiDua]

Summary

Training data construction is split by ownership: per-sample token fields are now built at generation time inside AgentLoop, while step-level conversion (validation, advantage, shift, tensorization, seq_ctx/teacher fields, data_info) moves into TrainingController. BaseRLTrainer._prepare_train_data is deleted.

• AgentLoop.canonicalize_train_fields (base + localhost/sandbox overrides) builds the unified full-sequence input_ids/labels/logprobs convention; semantic holes (tool/env tokens) are baked into labels as -100 by the loops themselves.
• RolloutState drops response_mask; labels becomes the only supervision carrier. New agent_loop_type records the producing loop class name (type(self).__name__), and AGENTIC_AGENT_LOOP_TYPES discriminates agentic full-sequence samples (None/unregistered types fall back to prompt+response semantics).
• TrainingController.fit accepts list[list[RolloutState]] and absorbs group validation, session-clustered advantages, the one-position shift, tensorization, seq_ctx/teacher fields and data_info statistics.
• calculate_group_effective_response_masks now bakes token staleness into labels in place (clearing is monotone, so replay-buffer expiry checks and the train-batch bake converge to identical labels); the agentic exclusion condition now keys on agent_loop_type instead of input_ids/labels presence, which would have silently disabled token staleness after canonicalization.

Behavior changes

• prompt+response samples whose field construction fails are marked FAILED at generation time (their group is skipped) instead of asserting at training time.
• Missing reward with task_adv_weight == 0 trains with zero advantage; with weight > 0 it still raises (crash semantics preserved, location moved).
• The prepare_data timing key disappears; conversion cost is now inside the training timer.
• prompt_len/response_len stats keep exact denominators for reasoning samples; only mixed-form edge cases may drift by +/-1 (log metrics only).

ProduceBatchResult impact: leftover_*/raw_rewards*/produced_* counts unchanged; generation-time canonicalize failures can flip individual samples to FAILED, which converts to group-level skips via validation (affects rollout_states content, not the accounting).
RoutedExperts impact: not affected; routed_experts is still attached to seq_ctx by the controller, ownership and release paths unchanged.
Ray concurrency impact: not affected; no agent-loop actor methods, decorators, or concurrency groups changed.

Test plan

• tests/rl/test_prepare_train_data.py rewritten as contract tests for canonicalize_train_fields (base + both agentic overrides) and TrainingController._convert_rollout_groups, including the packed-alignment regression (advantage used to be one position longer than input_ids).
• Staleness, replay-buffer, producer, multi-task manager and trainer tests updated to the labels-only contract; added regression cases: repeated staleness bakes converge, reset_rollout_response clears canonical train fields, canonicalized prompt+response states remain staleness-eligible.
• pre-commit (ruff, ruff-format, docformatter, pyupgrade, mypy strict) passes on all touched files; unit tests run in CI (pytest tests/rl).

[YanhuiDua]

@claude review

[github-actions]

Claude:

Summary

本次重构将训练数据构造按归属拆分：per-sample 全序列 token 字段（input_ids/labels/logprobs）在生成期由 AgentLoop.canonicalize_train_fields（基类 + localhost/sandbox/gsm8k/verl 各自覆写）构造，step 级转换（组校验、session 聚类 advantage、shift、张量化、seq_ctx/teacher 字段、data_info）下沉到 TrainingController；RolloutState.response_mask 被删除，labels 成为唯一监督载体，token staleness 直接烙进 labels。抽象与线性业务流程组织良好，未发现 Warning 及以上问题。

核心原理实现与单测

• Loop 侧 canonicalize：基类为 prompt+response 型样本装配全序列并对失败样本置 FAILED（隔离到单样本、不上抛 producer）；agentic 覆写只做统一约定校验。核心行为由 tests/rl/test_prepare_train_data.py 的 TestAgentLoopCanonicalizeTrainFields/TestAgenticLoopCanonicalizeTrainFields 经 public canonicalize_train_fields 覆盖（含 Tensor flatten、tokenizer 兜底、VLM train_prompt_ids、失败隔离），mock 仅限项目外依赖（lagent、tokenizer）。
• Controller 侧转换：_rollout_groups_to_colate_items 以清晰的 Phase 1-4 线性流程组织，advantage/shift/张量对齐由 TestConvertRolloutGroups 覆盖，并含 packed 对齐回归（advantage 曾比 input_ids 长 1）；test_rl_colocate_trainer_integration.py 经 public TrainingController.fit(rollout_groups, ...) 覆盖端到端真实路径。
• Staleness 烙制：calculate_group_effective_response_masks 原位清 -100（单调收敛），agentic 全序列样本首次纳入 token 级 staleness，由 test_multi_task_agent_loop_manager.py 与 test_staleness_policy.py 经 public produce_batch/calculate_group_effective_response_masks 覆盖（含全序列逐 token 烙制与同步路径零操作）。

其他 Issues

未发现 Warning 及以上问题。