[recipe, tests, doc] feat: add BAGEL FlowGRPO V1 sync recipe

[OnePunchMonk]

What does this PR do?

Claims the BAGEL + FlowGRPO (sync) row from #389's "Need help" table:
mechanical v0→v1 migration of the PickScore LoRA recipe, following the exact
pattern already used for SD3.5 (#296) and Wan2.2 (#466) — swap
main_diffusion → main_diffusion_v1 and add trainer.use_v1=true trainer.v1.trainer_mode=sync. No adapter/loss changes.

Changes

• examples/flowgrpo_trainer/bagel/run_bagel_pickscore_lora_v1.sh — v1 sync
  counterpart of run_bagel_pickscore_lora.sh.

• tests/special_e2e/run_flowgrpo_bagel_pickscore_v1.sh — tiny-random GPU
  smoke test, mirrors run_flowgrpo_bagel_pickscore.sh (which validates the
  v0 path), swapped to the v1 entrypoint.

• README: documents the new V1 recipe, cross-references

  docs/start/diffusion_v1.md.

Duplicate-work check

• gh pr list --repo verl-project/verl-omni --state open --search "bagel" —
  no existing PR touches BAGEL's v1 path (hits are an unrelated BAGEL SFT PR
  and an unrelated HunyuanImage-3.0 PR).
• [RFC] Tracking v1 upgrade for verl-omni: performance and compatibility #389's own dashboard (refreshed by a maintainer today, 2026-09-01) still
  lists BAGEL FlowGRPO under "Need help".
• Claimed in a comment on [RFC] Tracking v1 upgrade for verl-omni: performance and compatibility #389 before starting.

Checklist Before Starting

• Search for similar PRs. Query: gh pr list --repo verl-project/verl-omni --state open --search "bagel"
• Format the PR title as [{modules}] {type}: {description}

Test plan

GPU e2e (tiny-random smoke, validates the v1 entrypoint end-to-end):

bash tests/special_e2e/run_flowgrpo_bagel_pickscore_v1.sh

v0 vs v1, same container, 4x A10G (both scripts run back to back in a
single Modal GPU allocation, 8 steps each, tiny-random BAGEL + dummy
PickScore data): both passed with matching training/reward trajectories.

step	v0 actor/loss	v1 actor/loss	v0 reward mean	v1 reward mean
1	9.19e-06	4.30e-05	0.019565	0.019565
2	2.80e-05	3.08e-05	0.020356	0.020367
3	1.11e-05	2.89e-05	0.019334	0.019306
4	2.33e-06	1.31e-05	0.018826	0.018811
5	6.40e-06	2.35e-06	0.019532	0.019538
6	5.70e-06	1.46e-05	0.019638	0.019653
7	8.15e-06	1.31e-05	0.020528	0.020548
8	1.71e-06	2.28e-05	0.019014	0.019008

Both runs returned exit code 0. Reward means track each other step for
step (same seeded data), confirming the v1 sync migration reproduces v0
training behavior. Loss curves differ slightly (v1's actor loss runs an
early-run pass before the first weight sync, per the v1 trainer's
warmup), same order of magnitude, same downward trend, no divergence.

This was found and fixed in diffusers_training_adapter.py (see below and
the earlier PR comment) before this comparison could run at all — v1 was
crashing with a KeyError on every attempt until that fix landed.

AI usage disclosure

Claude Code was used to draft this migration (diffed against the already-
merged SD3.5 and Qwen-Image v1 recipes to confirm the exact mechanical
pattern) and the smoke test. A human submitter (onepunchmonk) has reviewed
every changed line and the Modal GPU run before taking this out
of draft.

Claude Code also found and fixed a real bug while testing the v1 path:
_prompt_token_ids_to_batch in diffusers_training_adapter.py read
micro_batch["prompt_token_ids"], a key nothing ever writes (every writer
of this field uses "prompts"). This crashed the v1 sync path with a
KeyError at the first rollout call. Fixed to read micro_batch["prompts"]
instead; this was required to unblock v1 testing for this PR.

[Copilot]

🟢 Approval recommended

Changes are isolated to new recipe/test scripts and README documentation, and the new v1 e2e script cleanly mirrors the existing v0 BAGEL smoke test with the expected mechanical v1 trainer switch.

Pull request overview

Adds a BAGEL + FlowGRPO v1 trainer (sync) recipe and a matching tiny-random GPU smoke test, following the established “v0 → v1 mechanical migration” pattern used by existing diffusion v1 recipes (switch to main_diffusion_v1 and set trainer.use_v1=true, trainer.v1.trainer_mode=sync). Also documents the new v1 entrypoint in the BAGEL recipe README and links to the diffusion v1 guide.

Changes:

• Add a v1 sync BAGEL PickScore LoRA recipe script using python -m verl_omni.trainer.main_diffusion_v1 with v1 trainer flags.
• Add a tiny-random GPU e2e smoke test script for the BAGEL v1 path, mirroring the existing BAGEL v0 e2e test but routed through the v1 entrypoint.
• Update the BAGEL recipe README to document how to run the new v1 sync recipe and reference the diffusion v1 guide.

File summaries

File	Description
tests/special_e2e/run_flowgrpo_bagel_pickscore_v1.sh	Adds a v1 “tiny-random” GPU smoke test for BAGEL FlowGRPO PickScore using main_diffusion_v1 + trainer.use_v1=true + trainer.v1.trainer_mode=sync.
examples/flowgrpo_trainer/bagel/run_bagel_pickscore_lora_v1.sh	Adds the v1 sync counterpart of the BAGEL PickScore LoRA FlowGRPO recipe script (mechanical v0→v1 entrypoint + flags).
examples/flowgrpo_trainer/bagel/README.md	Documents the new v1 sync recipe and links to docs/start/diffusion_v1.md.

Review details

• Files reviewed: 3/3 changed files
• Comments generated: 0
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[OnePunchMonk]

Ran the GPU smoke tests on Modal (4x A10G) to validate this against the v0 baseline.

v0 control (tests/special_e2e/run_flowgrpo_bagel_pickscore.sh): passed.

step:1 actor/loss=1.16e-05 critic/rewards/mean=0.0196
step:2 actor/loss=5.94e-06 critic/rewards/mean=0.0204
FlowGRPO BAGEL PickScore LoRA e2e test passed (training completed successfully)

v1 sync (tests/special_e2e/run_flowgrpo_bagel_pickscore_v1.sh): currently fails with

KeyError: 'key "prompt_token_ids" not found in TensorDict with keys [..., 'prompts', 'raw_prompt', ...]'

at verl_omni/pipelines/bagel_flow_grpo/diffusers_training_adapter.py:119, right at the first rollout call of step 1. Traced it to verl_omni/agent_loop/diffusion_agent_loop_tq.py, which writes the tokenized prompt to TransferQueue under the key prompts (matching v0), but nothing downstream renames/aliases prompts to prompt_token_ids for the TQ path the way the v0 trainer does before calling infer_actor_batch. So the BAGEL adapter's hardcoded micro_batch["prompt_token_ids"] lookup never finds it.

Also hit two environment-only issues along the way (not related to this PR, noting in case others hit them on Modal/similar sandboxes): a Ray placement-group deadlock caused by TransferQueue's default SimpleStorage.num_data_storage_units=8 competing for CPU with the actor bundle on smaller nodes, and a missing libGL.so.1 needed by cv2 (opencv-python via vllm-omni) on headless CUDA base images.

Will patch the prompt_token_ids gap and rerun both to build a v0-vs-v1 comparison.

[OnePunchMonk]

Pushed the prompt_token_ids -> prompts fix in diffusers_training_adapter.py (commit 895135b).

Kicked off a fresh v0-vs-v1 side-by-side run on Modal to confirm the fix and get comparable step metrics. Both jobs got killed by a Modal client disconnect during Ray init, before either reached a training step, so I don't have new numbers to post yet. Not a code issue, just an infra hiccup on this run.

Existing v0 baseline (from the earlier comment, still valid):

step:1 actor/loss=1.16e-05 critic/rewards/mean=0.0196
step:2 actor/loss=5.94e-06 critic/rewards/mean=0.0204
FlowGRPO BAGEL PickScore LoRA e2e test passed (training completed successfully)

Will rerun the v1 path and post updated metrics (plus a v0-vs-v1 plot) before this comes out of draft.

[knlnguyen1802]

cc: @zhtmike Do you think we need more step for verification

[zhtmike]

@AndyZhou952 do you have the last bagel record?

[AndyZhou952]

@AndyZhou952 do you have the last bagel record?

https://wandb.ai/andyzhou/bagel FYI

[zhtmike]

cc: @zhtmike Do you think we need more step for verification

the reward should not be such low

[OnePunchMonk]

@knlnguyen1802 @zhtmike flagging for visibility: the v0-vs-v1 comparison table above is a tiny-random smoke test (randomly-initialized BAGEL checkpoint + dummy PickScore data, 8 steps), not a full-checkpoint training run. It confirms the v1 sync migration is mechanically correct and reproduces v0's trajectory, but the flat ~0.02 reward is expected for random weights on dummy data, not a signal about real PickScore reward quality. A full-checkpoint validation would need a separate run.

[zhtmike]

I will tackle the validation for this as requested by @knlnguyen1802 .

And if you are interested, you may help to debug this as well: #561

[OnePunchMonk]

Sure, will check #561 next!