【Task.38】VAGEN-Lite：Sokoban 多轮视觉 RL 复现进度

[yuki-younai]

任务信息

• 任务：Task 38 VAGEN-lite（高级，论文复现类，NeurIPS 2025，arXiv:2510.16907）
• 目标：在 Relax 中复现 Sokoban/FrozenLake 多轮视觉交互 RL，支持 PPO/GRPO 与 concat/non-concat 两种训练模式，效果与官方 Lite 实现一致（成功率差 ≤3–5pp）。
• 参考实现：Relax-Task/VAGEN（基线 commit 2936322）
• 分支：feat/vagen-lite-sokoban（基于 task38-baseline）

当前进度

✅ 已完成

阶段	内容	验证
Phase 1 环境移植	Sokoban/FrozenLake env 迁入 examples/vagen_lite/（parse_response/prompt/seeding/patch_sokoban_env 逐字保留，仅改周边接线：async→sync、reward 走 info、build_env 派发）	12 单测通过：同 seed 同地图、同动作序列同 state/reward/done
Phase 1.5 多轮 harness	框架外隔离验证整条链路（env reset→VLM 生成→parse→step→下一 obs）	Qwen2.5-VL-3B 实跑，wm 模式 2 轮解出 Sokoban（reward=1.2）
Phase 2 Relax concat rollout	多轮循环搬进 generate(args,sample,sampling_params)->Sample 契约；env-driven 初始 prompt、reward=Σinfo["reward"]；复用 deepeyes helpers	隔离 smoke + 全 pipeline：Megatron 加载 + SGLang + rollout 实跑通（64 sample，含 success）
冷启动优化	/dev/shm warmup + --warm-hf-checkpoint-page-cache	Actor 加载 17min→35s，SGLang shard 11min hang→1s

🔧 train step 闭环（刚定位的两个 blocker）

rollout 已通，但 train forward 连续踩两个坑，均已定位修复，待重跑确认：

1. batch 整除约束：global-batch-size 8 → data.py 抛 mini_rollout_batch_size must be divisible by dp_size。改为 32（32/8=4，4%4=0；同时对齐 VAGEN-Lite train_batch_size=32）。
2. VL 多模态透传（根因）：actor forward 算 mrope 时 image_grid_thw=None → get_rope_index 崩。定位到 relax/utils/training/data_fields.py:18——multimodal_train_inputs 仅在 --multimodal-keys 非空时才透传到 actor。VAGEN 图来自 env（非数据集列）故未设此 flag。对齐 deepeyes_r3 加 --multimodal-keys '{"image":"images"}'：数据集无此列时 data.get("images") or [] 安全返回空，env-driven generate 会覆盖为真实 per-turn grids。

待办

• 重跑确认 train step 闭环（rollout→convert→advantage→forward/backward）
• Phase 3 non-concat 模式 + no_concat_gae advantage（需改 core：relax/components/advantages.py 新增分支 + turn 索引透传，按 CLAUDE.md「Ask First」征询）
• Phase 4-5 reward_variance 低方差 filter、FrozenLake、评测、401 步训练、与原版 Lite 成功率对比（验收 4）、docs/Issue/PR

验收对照

验收标准	状态
同 seed 同地图；同动作序列同 state/reward/done	✅
concat/non-concat 逐轮回放对齐	concat ✅ / non-concat 待 Phase 3
PPO+GRPO 各完成训练+评测无 hang/泄漏	GRPO concat train 待重跑确认；PPO 待 Phase 3
与原版 Lite 成功率差 ≤3–5pp	待 401 步对比

关联

• 详细规划：task-38-plan.md
• 进度记录：Relax/TASK38_PROGRESS.md
• 认领人：@yuki-younai

[KEY-ZZZ]

Hi @yuki-younai and maintainers, I'm @KEY-ZZZ. I'd like to join Task 38 without duplicating the work already in progress.

I can own a bounded, independently reviewable reproduction/evaluation slice using up to 4x80GB GPUs after the current train-step closure: deterministic parity checks, PPO/GRPO runs, and reporting for task success rate, throughput, peak VRAM, and training stability across concat/non-concat as those interfaces land. I can also take FrozenLake/evaluation if that is the higher-priority gap.

Could you confirm whether collaboration is still open and which files/metrics you would prefer me to own? I'll align with the current branch and plan before coding, and I won't present the work as an official Relax contribution until the scope is agreed.

[Michael-Salon]

Hello @yuki-younai，先确认一下：这个 task 有没有 RFC？方便的话可以贴一份出来，可以在未关联具体PR的情况下更好的review设计实现。

在 RFC 之外，我先基于 feat/vagen-lite-sokoban（fccf0c9，base task38-baseline @ e2325a0）这个分支上已经落地的 Phase 1-2（concat 模式）代码给一版评审意见：

---

1. 当前有 1 个公式实现的不匹配：

   • VAGEN-lite 官方说明是"built on VERL agent-loop"，而 VERL 的 compute_gae_advantage_return（verl/trainer/ppo/core_algos.py）在递推时显式跳过被 mask 的位置——nextvalues = values[:,t]*response_mask[:,t] + (1-response_mask[:,t])*nextvalues，lastgaelam 同理——即 observation/padding token 不消耗自己的 R_t/V_t，只是把上一个（未来）时间步的值原样传递过去，这与论文"Vanilla PPO fails due to lack of observation token masking"的诊断一致。当前 Relax advantage_estimator=ppo 分支（relax/components/advantages.py:180-192，以及本 PR 走 --colocate 时实际执行的同逻辑副本 relax/backends/megatron/loss.py:575-593）没有等价的屏蔽：values/kl/rewards 三个数组按 response_length 逐 token 排列，get_advantages_and_returns_batch（relax/utils/training/ppo_utils.py:584）函数签名不接收 loss_masks，observation token（含图像 token）和 action token 在同一条 γ/λ 递推链上一视同仁。对照组：reinforce_plus_plus/reinforce_plus_plus_baseline 都显式接收并用了 loss_masks，唯独 ppo 漏了。
   • 严重性限定：两个 PPO 脚本均为 --gamma 1.0 --lambd 1.0 --kl-coef 0.0。kl_coef==0 走 kl = zeros 短路分支，且 γ=λ=1 时 GAE 望远镜求和给出 A_t = ΣR − V_t，中间 observation token 的 value 完全抵消，masked 与 unmasked 在 action token 上输出一致（我按脚本超参做过数值对照，差值 2.4e-7；改成 γ=0.99/λ=0.95 后差值 0.356）。因此本轮跑分不受影响，这条属于潜在缺陷：一旦 gamma/lambd 偏离 1.0 或启用 KL 惩罚，就会静默偏离参考实现。
   • 说明：issue 里提到的参考实现 Relax-Task/VAGEN（commit 2936322）我这边访问不到（私有仓库/404），以上比对基于公开的 RAGEN-AI/VAGEN 及其依赖的 verl fork JamesKrW/verl@vagen-lite（masked GAE 见该 fork core_algos.py:252-253）。若内部参考实现与公开版有出入，还请作者或有权限的同学指正。

2. 风险告知：

   • advantages.py:188 的 k[-1] += reward 把整条轨迹的累计 reward 加到 response_length 数组的最后一个 token。但 rollout_concat.py:207-220 里，轨迹以 budget_exhausted/max_turns 结束时（没能在 max_turns 内 done=True 的常见情况），最后一个 token 是刚 append 的 observation（loss_mask=0）。参考实现在这一点上的处理不同：VAGEN 的 concat agent loop（vagen/agent_loop/gym_agent_loop.py:350-382）先判终止再决定是否 append observation，done/turn-limit/token-limit 三条均在 append 之前返回 TERMINATED，docstring 明确写为 "stop WITHOUT appending user suffix, so the episode ends on an assistant turn"，配合 agent_loop.py:575 的 rm_scores[:, response_length] = scores，其 reward 始终落在 response_mask=1 的 assistant token 上。需要说明的是，Relax 当前不做 masking，GAE 反向递推会把末位 reward 传递给全部在先 token（γλ=1 时无衰减），因此 outcome reward 并未丢失；但该落点与 finding① 存在耦合——verl 的 masked 递推在 mask=0 位置不累加 delta，若仅补 masking 而不调整 rollout 顺序，这类轨迹的 reward 将被真正丢弃且不报错。两项改动的先后不可颠倒。
   • observation token 参与 GAE 递推本身可能污染相邻 action token 的 advantage：compute_approx_kl 对整条序列逐 token 计算，不区分 loss_mask，图像/环境文本 token 上策略与参考模型分布可能差异很大，KL 异常值经 γ·λ 衰减渗透进相邻 response token，方向和幅度不可控。该风险在当前两个 PPO 脚本上不成立（--kl-coef 0.0，compute_approx_kl 不会被调用），仅在后续启用 KL 惩罚时需要关注。
   • 上述两条风险目前确认只影响 advantage_estimator=ppo（GAE + critic）路径；grpo/gspo/sapo/cispo 走常数广播（get_grpo_returns），不受影响。但验收清单里 PPO 和 GRPO 都是要跑的训练模式，run_vagen_lite_sokoban_ppo_concat.sh / run_vagen_lite_frozenlake_ppo_concat.sh 都在范围内。

3. 实现建议：

   • 短期：建议将 rollout_concat.py 的 budget_exhausted/max_turns 判定移至 _append_to_sample(..., loss_mask_val=0) 之前，对齐参考实现的"episode ends on an assistant turn"。这是 examples/ 内的改动，不触碰 core，不需要走 CLAUDE.md「Ask First」流程；改完之后 reward 落点即与参考实现一致，masked/unmasked GAE 在任意 gamma/lambd 下也都对齐。基于 finding① 的等价性结论，PPO 在当前超参下可以直接参与验收标准④的成功率对比，无需降级为"仅供参考"。
   • 不建议在通用的 relax/components/advantages.py 的 ppo 分支里为 VAGEN 单开 if-else 塞 masking 逻辑（会牵连其它走 ppo 估计器的 recipe，比如 deepeyes 的多轮 tool-call 场景）。如果要修，建议参照 reinforce_plus_plus 已有模式，给 get_advantages_and_returns_batch 加一个可选 loss_masks 参数，并同步修改 loss.py 中的同逻辑副本，这是框架级改动，需要按 CLAUDE.md「Ask First」流程走，优先级排在上一条之后。
   • 建议先补一条单测覆盖风险①（构造一个不会 done=True、跑到 max_turns 结束的轨迹，断言最后一个 response token 的 loss_mask==1，且 action token 的 advantage 中包含完整的轨迹 reward）——现有 tests/examples/test_vagen_env_determinism.py 的 12 个测试都是环境确定性测试，没有覆盖到 reward→advantage 的落点问题。
   • 请标注哪些地方 Relax 与参考实现不一致：一是 rollout 终止顺序（参考实现终止于 assistant turn，Relax 会先追加 observation）；二是 ppo 估计器未做 observation token 的 GAE masking（当前超参下与参考实现数值等价，超参变更后会偏离）。GRPO 路径不受这两项影响。另外已核对一致、无需改动的两处：advantage 白化（脚本已带 --normalize-advantages，loss.py:685 走 distributed_masked_whiten(all_advs, loss_masks)），以及 policy/value loss 的 reducer 均按 batch["loss_masks"] 归约（loss.py:1272/1284）。