feat(agentic): add Sample-native multi-turn rollout and training - #214
Merged
Conversation
…-446) Squash merge of origin/LIN-446/main into LIN-444/main. Combines: - add Sample/Part endomorphism types - derive Sample/Part lineage from the sample_ids path
Replace the RolloutReq/RolloutResp request/response pair with the unified recursive Sample/Part types across the dedicated rollout engines — generate: Sample → Sample, in place, no bridge. The model pipelines, the trainside engine, the trainers, reward, and the train stack stay on the old types for a later pass, so this branch is intentionally not runnable as a full training loop yet. Engines converted: - sglang (AR text/vlm), sglang_diffusion (image), vllm_omni (sd3 / qwen_image / hv15 single-DiT + HI3 multi-stage), composed (PE AR→diffusion chain). - base.py ABC flipped to generate(self, sample: Sample) -> Sample; engine/__init__ chunked_engine_generate_req -> chunked_engine_generate (split-by-root -> regroup -> concat, the Sample analogue of slice-by-index). Mechanics: - Pre-forked gen shells are filled by sampling_params type; positional lineage replaces parent_track; σ pinned onto DiffusionSamplingParams.sigmas; the x_T noise key is derived from the path-id lineage (OD-2); stage_config moved to Part.control. - Multi-input multimodal (chained input Parts) is gated with clear deferral raises — a deferred non-goal of this pass. P0 type fields: Part.control / conditions / fill, DiffusionSamplingParams.sigmas / init_noise_latent_shape, LatentSegment.initial_latents. Verification (rollout-only, no training): - scripts/check_sample_roundtrip.py structural oracle (5/5 contracts). - Real GPU rollout smokes pass for every converted engine and modality class: AR (sglang/Qwen3), diffusion (vllm_omni sd3, sglang_diffusion sd3), and the multi-stage composed PE chain (AR→diffusion).
Wire the previously-gated image+text / cot_text modalities (hi3_it2i, hi3_i2t, the sglang vlm path, and hi3_dit_recaption) onto the generate(Sample) -> Sample boundary. The type layer already supported it, so this is mostly un-gating: a second input rides as a chained input Part via the new Part.input_child(primitive) (branch-1 child, no sampling_params) so only the head stays a root, and Sample.conditioning() surfaces both primitives in turn order. Adapters locate the chained input by primitive type (image_input_part / new cot_text_from_sample); hi3_dit_recaption gets the ported cot_text build. Also fixes sglang text.py build_response to preserve intermediate input Parts (return [*parts[:-1], filled] instead of [parts[0], filled]) so a multi-input chain's gen Part keeps a valid parent. Verified engine-only with hand-built Samples: oracle check_sample_roundtrip 7/7 (incl. multi-input image-chain + cot_text fixtures), a real Qwen2.5-VL VLM run, and a full HunyuanImage-3 80B IT2I GPU run (image+text -> AR recaption -> edited image, chain filled).
Migrate the consumer side (all five trainers + train stack + wandb logging) from RolloutReq/RolloutResp/RolloutTrack to the Sample/Part endomorphism model (LIN-446) — the counterpart to LIN-454's rollout-engine conversion. Trainers now build request Samples (Part.input + .fork) and consume response Samples (frontier = parts[-1], stages located by sampling_params type); GRPO advantage / reward-propagation / token-balancing run through the methods already on Part/Sample. Scope: all trainers incl. it2i (multi-input via Part.input_child) and the HI3 unified two-engine path. Out of scope / unchanged: the reward service (score_and_attach migrates on a separate branch — its (*, req, track) signature does not yet match the trainers' score_and_attach(sample) call, so the train tail is gated on that branch), the trainside engine, and the rollout-engine internals. Highlights: - Sample/Part navigation on the type: root_group_ids, gen_parts/gen_part/ gen_part_index, with_parts, and tree-aware slice/select (a request Sample DP-scatters by prompt-tree, not by the parts list). - Per-trainer request builders + the HI3 two-engine _run_rollout_one stitching (AR recaption -> DiT image); the DiT x_T noise key is re-rooted from the globally-unique lineage so it stays unique across dp>1 replicas. - UnifiedModelTrainStack.train_track scatters the [input, ar, image] lineage as a unit so both stages shard at the same prompt boundaries. - wandb on Sample/parts: diffusion stage named "diffusion" consistently; num_samples reports the generated-sample count; zero-std groups bucket by the advantage grouping (advantage_group_ids). - disable_driver_xt flag on DiffusionSamplingParams (restores the unified DiT driver-x_T escape hatch). Verification: scripts/check_sample_roundtrip.py (14 contracts) + SD3 (vllm-omni), AR (sglang), and HI3 two-engine generation smokes pass on GPU. The dp>1 cluster is contract-tested + A2 engine-verified; its multinode end-to-end behavior is not yet verified, and the full train loop awaits the separate reward migration.
Convert the model bundle's rollout boundary from RolloutReq/RolloutResp to the
recursive Sample/Part types — the model-bundle half of the refactor, pairing with
LIN-480 (trainers / train-stack):
- SD3Pipeline.generate / Qwen3Pipeline.generate take a request Sample (pre-forked
gen frontier) and fill it; a factored _conditions_for() is the single encode
path shared by rollout and trainer-side replay (re-encode; Part.conditions left
empty — no cache).
- TrainsideRolloutEngine.generate -> Sample -> Sample: _ensure_sample_sigmas pins
sigma onto the gen part's DiffusionSamplingParams; Part.slice/concat chunking.
- NoiseRecipe.from_sample() — Sample-shaped x_T builder keyed on the gen part's
path-id lineage.
- check_sample_roundtrip.py +2 contracts (from_sample lineage, generate-fills-
frontier); trainside_{sd3,ar}_smoke.py real-GPU rollout+replay smokes.
Validated on H20 (pod unirl-gz-2): CPU oracle green; SD3 rollout->replay
mean|dlogp|=0 (ratio~1); Qwen3 AR replay self-consistency 0 (the
old_logp_source=replay contract — the in-process autoregress bf16 record vs
replay fp32 differs ~3-5, harmless since trainside recipes use replay for
old_logp). Stages (diffusion.py/ar.py) untouched.
Convert the reward adapter from RolloutReq/RolloutTrack to the Sample/Part endomorphism model — the consumer-side counterpart LIN-480 left gated. The trainers already call score_and_attach(sample); this matches that contract. - RewardService.score_and_attach(*, req, track) -> RolloutTrack becomes score_and_attach(self, sample: Sample) -> Sample: score the frontier Part, build input primitives from Sample.conditioning() (nearest-ancestor caption), generated from frontier.primitive keyed by the backend preferred_input_kind, and root-sourced prompt metadata via the new Sample.root_metadata. - AR truncation/overlong shaping preserved verbatim, re-gated on the scored frontier being ARSamplingParams. - Delete the request/track expansion + metadata-normalization machinery (_KIND_TO_KEY, _normalize_prompt_metadata, _build_request_for_track): the Sample is row-aligned and DP-scatters by prompt-tree, nothing to reconcile. - Add primitive_modality_key + Sample.root_metadata; refresh reward docstrings/README to Sample/Part. - Add scripts/check_reward_roundtrip.py structural oracle (6 contracts; passes on-GPU-pod with real torch).
…n-stack replay SD3Pipeline.generate left Part.conditions empty (a "replay re-encodes" assumption), but the train stack's prepare_segment reads part.conditions directly with no re-encode wiring — so trainside SD3 training crashed in SD3Conditions.from_dict (text=None). Every sibling diffusion pipeline and the SD3 sglang_diffusion adapter already populate conditions via fill(conditions=...); bring trainside SD3 in line. Surfaced by the SD3 trainside e2e (trainside_sd3_smoke bypassed this by calling stage.replay directly, never going through train_track).
…rollout texts_from_sample asserted prompt count == gen sample count, but the request keeps the input Part un-fanned (one prompt per group) while the gen shell fans out samples_per_prompt via Sample.fork — so SD3/qwen-image served rollout crashed on rollout 1 with "prompt count 6 != gen sample count 96". vllm_omni runs num_outputs_per_prompt=1 (one output per gen sample), so tile each prompt across its group-by-parent-contiguous gen siblings. A 1:1 request (no fan-out) is unchanged. Surfaced by the SD3 sd3_vllmomni e2e (served LIN-454 path).
…feUnpickler sglang 0.5.12's SafeUnpickler (CVE-2025-10164 guard) runs in the engine process and blocked the colocate tensor weight-sync: update_weights_from_tensor deserialization rejected unirl's vendored FlattenedTensorBucket / _rebuild_cuda_tensor_modified (both under the unirl. namespace) with "Blocked unsafe class loading". Register "unirl." on sglang's server-side allowlist in SGLangRolloutEngine.__init__ (the tp_size=1/use_ray=False scheduler runs in this process), mirroring what unirl's own SafeUnpickler already permits. Surfaced by the Qwen3 DRPO AR e2e (qwen3_drpo_4b_base_dapo_sglang): the full vertical (sglang gen -> MathVerify -> DRPO -> FSDP train) ran and crashed only at the first weight-sync.
… compat) The prior allowlist patch (a1f33d1) registered "unirl." on the engine process's SafeUnpickler, but sglang's SRT scheduler deserializes the weight bucket in a SEPARATE spawned subprocess with a fresh allowlist, so the patch never took (same "Blocked unsafe class loading" recurred). Instead, serialize with sglang's NATIVE FlattenedTensorBucket / MultiprocessingSerializer / monkey_patch_torch_reductions (verified identical device-UUID IPC mapping) so the wire references sglang's own allowlisted classes — matching what the sglang_diffusion engine already does. Reverts the ineffective __init__ allowlist patch. Surfaced by the Qwen3 DRPO AR e2e: full vertical ran, crashed only at the initial colocate weight-sync into SRT.
…hape
Convert every model pipeline's generate() from RolloutReq → RolloutResp to the
recursive Sample/Part endomorphism, completing the model-bundle half of the
rollout refactor (pairs with the earlier LIN-454/479/480/481 engine/trainer/
train-stack/reward migrations), then retire the legacy types.
- Base contract: Pipeline.generate(self, sample: Sample) -> Sample; lift the
request-side Sample readers out of the rollout layer into types/sample_ops.py
so pipelines import them without a models→rollout inversion.
- 7 Tier-1 diffusion pipelines (qwen_image, z_image, wan21, wan22, hunyuan_video,
hunyuan_video15, ltx2): read prompt via sample.conditioning(), gen params off
the frontier Part, fill the frontier shell; factor _conditions_for() shared by
rollout + trainer-side replay. i2v image arrives via the input_child chain.
- qwen_vl + flux2_klein (AR/edit with one extra input modality).
- pe, bagel, hunyuan_image3 (multi-stage 2-gen-part flows): PEPipeline consumes
the pre-forked [input, ar_shell, diff_shell] Sample and drives the child
pipelines (mirrors the migrated ComposedRolloutEngine); bagel dispatcher reads
task from parts[0].control and fills one or both gen Parts; HI3 dispatcher + 5
modes/ (incl. two-part t2ti).
- Retire unirl/types/rollout_req.py + rollout_resp.py; relocate PrimitiveValue to
types/primitives.py; drop the dead NoiseRecipe.from_rollout_req; sweep doc
breadcrumbs across conditions/algorithms/sde/engine docs.
All pipelines carry conditions=<conds>.to_dict() on frontier.fill — the train
stack reads Part.conditions (GRPO/FlowGRPO re-type via conditions_cls.from_dict);
it does NOT re-encode. (Caught by E2E pe-trainside on unirl-gz-2: an earlier
revision left conditions empty and crashed at the first AR train step with
"Qwen3ARConditions.from_dict: expected d['prompt'] ... got None".)
Validated on TaiJi unirl-gz-2 (8×H20, GZ): compileall green; sd3-trainside reward
grew 0.74→0.79 with ratio=1.0 / advantage_std≈0.49; pe-trainside (SD3+Qwen3-0.6B)
runs the full 3-part chain with both AR + diffusion stacks training.
Adds scripts/trainside_{qwen_image,qwen_vl,pe}_smoke.py (tier-representative GPU
rollout→replay smokes).
…ti-turn encode (both backends) Squash of LIN-503/main (8 commits) onto LIN-444/main. Takes the rollout Sample from single-turn-only to full multi-turn, multi-modal conditioning end-to-end. - Role-aware trajectory layer on Sample/Part (Phase 1/2): Turn, Part.role / resolved_role, turns(), text_conditioning() / vision_conditioning() fail-loud renderers, replace_frontier / with_filled_frontier write-back; diffusion/omni engines wired; sample_ops retired. - gap C — sglang AR engine encode consumes the trajectory: text.py / vlm.py build the chat conversation from text_/vision_conditioning() via a pure transpose + de-expand util (rollout/engine/sglang/utils/conversations.py); resolve_sampling fan-out fixed to the last-fork branch (parts[-2]). - Trainside conjugate — qwen3 / qwen_vl in-process encoders consume turns() via models/types/conversations.py (transpose, no de-expand, inline-PIL VLM fusion); stale "_conditions_for re-tokenize" docstrings + trainside smokes fixed. - Behavior-preserving on single-turn (byte-identical); parity-safe by construction (replay teacher-forces over the stored conditions["prompt"], not a re-encode). - Verification: CPU oracles (engine 8/8, trainside 6/6, layer 7/7) + 4 GPU multi-turn smokes (sglang/trainside × text/VLM) green on Qwen3-4B-Instruct and Qwen2.5-VL-3B-Instruct — the captured prompt carries user→assistant→tool (+ image), replay ratio=1.
…t (LIN-492) Add the agentic (multi-turn tool) rollout layer over the Sample/Part model, on top of LIN-503 multi-turn conditioning. Scope is ROLLOUT only — trainer wiring, reward, and a dataset for agentic GRPO are a follow-up (the trainer is still single-turn). - AgentLoop (unirl/rollout/loop/): environment-driven synchronous loop — fork -> generate -> env.step -> observe, bounded by max_turns. Structural Environment / RolloutEnginePort protocols; the existing SGLang engine satisfies the port unchanged. - ToolEnvironment + tools/ (Tool, CalculatorTool): parse <tool_call>, run a registered tool, feed the result back as a role="tool" observation, stop on a final answer (or max_turns). Safe ast-based calculator (no eval). - Sample.observe(role="tool"): append the world-response as a mask-0 input Part tagged for LIN-503 role-aware rendering (<tool_response>). Additive — no existing code path changed; the whole merge is 1648 insertions, 0 deletions. Verified: - CPU oracles: scripts/tool_env_smoke.py (9 checks), scripts/rollout_loop_smoke.py (6), scripts/check_sample_roundtrip.py (16, regression) — all green. - GPU (Qwen3-4B-Instruct): single-turn tool call (tool_env_ar_smoke); the closed loop — model calls the calculator, sees the exact result, and answers (tool_env_ar_loop_smoke); and a 72-trajectory reliability run at 97.2% correct with clean termination (tool_env_ar_reliability). Known follow-ups: trainer integration (multi-gen-Part train step + reward + dataset), per-sample termination for n>1 heterogeneous batches, more tools.
… agentic) Squash-merge of LIN-499/main: the async per-group engine contract (agenerate core + sync generate facade + abort/pause/resume), all five engines migrated, Part.weight_version provenance, DevicePool.worker_max_concurrency knob, and the CPU async-contract test suite (11 tests pass on the merged tree). Conflict in types/sample.py resolved by keeping BOTH agentic's Part.role (LIN-503) and LIN-499's Part.weight_version — adjacent new fields, not mutually exclusive.
…ator) AgenticRolloutEngine drives multi-turn (tool-use) rollout: a rank-0 coordinator (BROADCAST+RANK_ZERO, the NCCLWeightSync pattern) over a DP slab of per-worker persistent drain loops. Each worker runs one run_until_complete that pulls single-trajectory tasks from rank 0 and runs them as multi-turn agent loops on its inner engine's event loop (continuous-batched via the inner backend's semaphore). generate returns a flat List[Sample] of variable-depth trajectories; the GRPO n-group is recovered by bucketing on the shared prompt root id. - base: widen generate to RolloutOutput = Sample | List[Sample] (single-turn returns Sample, agentic returns List[Sample]); per-turn seams stay -> Sample. - engine/agentic: coordinator (set_workers/generate/next_task) + per-worker drain (run_drain/_drain/_run_one/_pull); _run_coro delegates to the inner so the drain and weight-sync verbs share one lock (the quiesce boundary); a per-worker trajectory cap distinct from the backend request semaphore; lifecycle/weight-sync verbs delegate to the inner engine. - loop: add async astep to the Environment protocol + ToolEnvironment (non-blocking tool boundary); make step re-entrant (turn derived from the sample, not a mutable counter) so one env serves concurrent trajectories; add agenerate to RolloutEnginePort. - tests: CPU contract tests (ragged List[Sample] + bucket-by-root, cap saturation, two distinct bounds, pull load-balancing, failure isolation, FIFO next_task) + astep parity / re-entrancy / slow-tool-yields. - scripts: multi-worker GPU smoke. Validated on a 2-worker Qwen3-8B H20 slab: ragged List[Sample] with correct tool-use answers, and cross-worker rank-0 aggregation proven safe (gen segments produced on 2 distinct workers, all spans plasma-backed in Ray's global object store, 48 TensorRefs hydrate from the driver) — risk #1 closed, no run_drain materialization needed. LIN-522
…earch tools, ALFWorld Land the full multi-turn agentic RL capability on top of AgenticRolloutEngine. AgenticTrainer (unirl/trainer/agentic.py) GRPO over variable-depth trajectory lists: group-relative advantage across the n siblings of a prompt, every assistant turn concatenated into ONE on-policy train_track step (ratio ~ 1). The per-trajectory reward step (_rewards_and_groups) is overridable so tasks swap only the reward SOURCE. Deep-research task (answer-graded) M1 calculator + MathVerify and M2 search/visit tools + LLM-judge reward, with recipes and the train_deep_research entrypoint. ALFWorld baseline (AgenticEnvTrainer, env-sourced reward) Stateful per-trajectory ALFWorld adapter, an engine hook that attaches the environment's terminal-success return to the last turn, admissible-action snapping to avoid TextWorld PDDL crashes, and crash-exclusion (NaN reward -> neutral, zero advantage) so env-bug crashes cannot corrupt the GRPO gradient. Verified on-pod: task-success reward rises ~0.53 -> ~0.85 (peak 0.94), on-policy, on a fixed 8-game set (genuinely multi-turn: up to 10 env steps per trajectory, n=8 GRPO group). Testing: CPU contract tests cover the plumbing (reward attach, group advantages, env adapter, train_step assembly). The GPU end-to-end path is validated empirically (the rising curve), not yet in CI.
Add a StatefulTool seam so tools can hold per-trajectory state across turns, with guaranteed teardown and a first out-of-process tool. - StatefulTool(Tool): session_start -> execute_session -> session_end, keyed by a session id carried in the root control bag. ToolEnvironment dispatches on isinstance, so the stateless Tool path is unchanged. - ToolEnvironment: reset mints a uuid4 session id per stateful tool and stamps it into control["tool_sessions"] via _part_with_field (returns the request unchanged when there are no stateful tools); step/_run dispatch execute_session; async aclose ends sessions in the executor, swallowing errors. - AgenticRolloutEngine._run_one: finally-hook calls env.aclose on every path (success, crash, abort), duck-typed via getattr and wrapped so it can never re-raise into the drain. - AlfworldEnv.aclose: reclaim the episode + pooled template — fixes a leak when a trajectory dies in the engine between turns. - SandboxTool: persistent per-session Python REPL subprocess (lazy spawn in the executor, select-based timeout, killed on session_end); validates cross-turn variable reuse. - Tests: StatefulTool lifecycle + teardown on success and forced exception, SandboxTool REPL, and an ALFWorld leak regression.
Over-generate, interrupt generation at a turn boundary on weight sync, commit the fast complete GRPO groups, carry/drop the slow tail. - Engine: submit/poll/abort/drained over a background buffered drain; turn-boundary checkpoint-and-resume in _run_one. - Colocate driver: AgenticPartialTrainer / AgenticEnvPartialTrainer — keeps all GPUs on generation while cutting the straggler tail. - Disaggregated driver: AsyncAgenticTrainer / AsyncAgenticEnvTrainer. - Tail policy: carry (resumable tool envs) vs drop (stateful envs, ALFWorld). - Shared _advantage_train_and_log extracted from the barrier train_step (barrier path preserved); entrypoints + ALFWorld/deep-research recipes; per-rollout turn-histogram + committed/dropped instrumentation. ALFWorld speed study: colocate-partial beats the barrier only with group-level depth variance — Qwen3-8B ~18% faster, Qwen3-0.6B 27% slower. Integrates with LIN-533 teardown: the _run_one finally aclose now also fires on the abort/checkpoint path, releasing carried trajectories' episodes/sessions.
…udge) Reproduces AReaL's tongyi_deepresearch deep-research agent in UniRL's agentic stack (sync AgenticTrainer + AgenticRolloutEngine): - Verbatim Tongyi SYSTEM_PROMPT for the M2 recipe (deep_research_search_judge) - Hardened SearchTool: serper + serpapi providers, retries/backoff - Hardened VisitTool toward AReaL tool_visit.py: Jina read retries, structured evidence/summary extraction, content truncation - Robust LLM-judge verdict parsing (negative-first regex; fixes "not correct") - Cross-config invariant guards (AgenticTrainer + AgenticRolloutEngine) - prepare_asearcher: default ASearcherBase35k split, streaming load - Committed GPU training smoke (scripts/train_deep_research_smoke.py) Validated live: reward climbed 0.22 -> 0.50 over 19 rollouts (Qwen3-1.7B policy, Qwen2.5-72B judge, batch 128, n=8) before an infra-driven collapse. Follow-ups (not in this change): multi-turn token-recording reconciliation (Miles-style trim/accumulate) and a trajectory context-token cap for long-horizon runs. The shipped recipe defaults to max_turns=8 (safe).
CPU unit test for _parse_verdict guarding the negative-first parse against the "incorrect" ⊃ "correct" substring trap and "not correct" / "wrong" phrasings the prior substring test misread. 20 cases; the coverage Phase C intended but the M2 squash omitted. Verified: 20 passed on agentic@ce13a48d.
Collapse the grouping/lineage duplication on the Sample/Part model: - sample_id: add ancestor_id(sid, depth) — grouping labels are id-prefix projections. root_group_ids and Sample.split project directly; the _root_groups_per_part walk (re-deriving what __post_init__ validated) is deleted. - compute_advantages: the group_ids label-list override becomes group_layer (None = immediate parent, 0 = root prompt; PE diffusion_group_scope="prompt" now passes group_layer=0 instead of threading root_group_ids labels). scope stays the normalization mode, with the historical global branch (unbiased std) kept verbatim for bit-parity with the shipped adv_normalization_scope: global recipes. - delete Part.advantage_group_ids (its only reader, the zero-std wandb metrics, buckets by sibling group_ids) and the dead Part.split. - kind signal: fork requires sampling_params (a paramless gen shell is unrepresentable); Part.is_gen is the single predicate behind gen_parts / resolved_role / the wandb gen filter. Ids stay pure addresses; role stays presentation-only. - contracts: check_group_layer_advantages pins per-layer GRPO values, the unbiased-std global formula, and fail-loud on paramless fork and negative group_layer.
…backends Drop agenerate/run_session/CoroutineFactory from the rollout-engine contract (BaseRolloutEngine, BaseSingleTurnRolloutEngine) and make the SGLang backends safely callable from concurrent threads, so the agentic drain can drive sync generate from one thread per trajectory: - native: replace the SessionRunner run_until_complete-under-drive-lock model with a serve/park LoopThread over SGLang's own engine.loop — callers submit coroutines threadsafe and stay in flight together; weight/memory verbs require quiesced generation and run with the loop parked (the Engine's sync wrappers still drive the idle loop themselves). - http: pure sync (urllib + threading.Semaphore + per-batch thread fan-out); the httpx client and the backend-owned event loop are gone. Controls are bounded 10s best-effort POSTs. - Backend protocol: generate is sync + thread-safe; async surface removed. New CPU tests cover the LoopThread lifecycle (concurrent overlap, semaphore bound, park/serve, quiesce guard, close-waits) and the sync HTTP backend against a stdlib ThreadingHTTPServer stub.
trainside, vllm_omni, sglang_diffusion, composed: delete agenerate (a to_thread wrapper over the locked sync path) and run_session + the LocalAsyncRuntime each held only to serve it; their sync generate paths are unchanged (the generate lock stays their concurrency story). Drop agenerate from the RolloutEnginePort protocol (its only consumer, AgentLoop, is sync). Delete unirl/rollout/engine/runtime.py and its test file; SessionRunner keeps a local transitional CoroutineFactory alias until the test fakes move off it.
…suite AgenticRolloutEngine now drives trajectories on a per-drive thread pool (per_worker_concurrency threads = the trajectory cap; a thread holds its trajectory across tool-wait, preserving the two-bounds design) instead of coroutines on the inner engine's loop: _drain_worker pulls (blocking ray.get) and runs a fully-sync _run_one (inner.generate per turn + env.step + duck-typed env.close teardown). run_drain joins every drain thread and re-raises the first failure; a failed worker sets _stopping so siblings checkpoint. The barrier agenerate (test-only) is deleted; coordinator verbs and the turn-boundary checkpoint/carry semantics are unchanged, and the weight-sync quiesce is now 'abort joined the drain threads => decode-idle'. Environment protocol: astep/aclose (bridges that existed only for the deleted loop model) are removed; sync close(sample) joins reset/step as the guaranteed teardown hook (ToolEnvironment._end_sessions / AlfworldEnv._teardown_episode bodies, now public). SessionRunner is deleted (backends/base.py is protocol- only again). Tests: _fakes.py is a sync thread-safe backend/engine pair with a hold-gate for deterministic overlap; the engine tests assert the same contracts under threads (cap saturation, two distinct bounds, checkpoint/ resume, conservation, continuous batching); loop tests use step/close, with thread-based re-entrancy replacing the astep suite.
Rewrite the README's generation-interface and extending sections for the sync-only contract (threads, not asyncio; concurrent-caller requirement for agentic inners; the loop survives only inside the native SGLang backend). Update tool/env docstrings that still referenced the deleted astep/aclose/ SessionRunner, and guard the HTTP batch path against an empty wire (the old gather path returned [] there; a zero-worker pool would raise).
Keep the Sample/Part architecture while porting main's model, rollout, reward, and training behavior. Migrate decoded outputs to typed primitive maps and remove retired rollout response APIs.
# Conflicts: # unirl/rollout/engine/sglang/utils/sampling.py
Collaborator
Author
|
PE: |
Collaborator
Author
|
qwen omni: |
Collaborator
Author
|
DRPO: |
Collaborator
Author
|
SD3 vLLM Omni: |
The six agentic entrypoints were the only ones in the repo named after a dataset. Every other entrypoint is named for its trainer -- trainer/ar.py -> train_ar.py, whose docstring states outright that it serves both qwen_vl and qwen3 recipes. Apply that existing convention: train_deep_research.py -> train_agentic.py train_partial_deep_research.py -> train_agentic_partial.py train_async_deep_research.py -> train_agentic_async.py train_alfworld.py -> train_agentic_env.py train_partial_alfworld.py -> train_agentic_env_partial.py train_async_alfworld.py -> train_agentic_env_async.py Each name now maps 1:1 onto its trainer module, and both axes are legible: _env is the reward source (environment return vs graded terminal answer), _partial/_async is the execution topology. Topology is a suffix so the whole family groups under train_agentic*; train_async_ar.py keeps its prefix rather than widen the diff. Recipes stay benchmark-named -- the entrypoint names the capability, the recipe names the task. Each docstring gains a train_ar.py-style line naming the recipe family it serves, and the README tables become reward-source x topology matrices instead of benchmark lists. Bodies are unchanged: same trainer classes, same constructor arguments, same config_name defaults. No behavioral change. Note: python -m unirl.train_alfworld and train_deep_research no longer exist; six user-facing module paths moved.
Three files added by adf70e8 were never run through ruff-format, so `pre-commit run --all-files` -- which the Lint workflow runs on every PR to main -- fails on this branch. Formatting only; no logic touched.
… barrier 7ad5b34 established that an infrastructure fault must be excluded from GRPO rather than scored as a genuine miss -- "scoring an infrastructure fault as a genuine miss manufactures a gradient for every sibling in the group" -- but it only reached AgenticEnvTrainer. The colocate-partial and fully-async env variants each carried their own copy of _rewards_and_groups and kept the pre-fix 0.0, agentic_env_async.py even retaining the comment "gen-less / failed trajectory stays a legit group member" that the fix says it reverses. Concretely: when a trajectory faults before its first turn the engine attaches no reward (_attach_env_reward returns early on an empty gen list), so the trainer's missing-reward branch decides the sentinel. NaN is dropped from the group's mean/std and given zero advantage; 0.0 instead drags the mean, shifts the std, and hands every sibling a gradient the model never earned. A fault AFTER turn one was already fine on all paths -- the engine attaches NaN there and it passes straight through. Rather than patch two literals, hoist the method into a shared _EnvRewardSource mixin that all three env trainers mix in ahead of their trainer base. One implementation, so the next sentinel change cannot land on some paths and miss others. Net -83/+30 in the trainers; the reward SOURCE becomes a named concept matching how the docstrings already describe it. Tests, restoring coverage this branch lost: - tests/rollout/test_agentic_failure_marking.py, deleted by c566479 "test: remove tests directory" while the fix it guarded stayed. Restored verbatim; it still passes against current code. - tests/trainer/test_agentic_env_reward_source.py, new. The gap that let this bug through: nothing ever asserted the env TRAINERS' sentinel. Parametrized over all three env paths, plus a structural test that they resolve to one shared function. Verified against the pre-fix tree: exactly 3 failures, naming the two buggy paths and the divergence itself. 18 pass after the fix.
Removes the 13 remaining test files (1117 lines), including the two added in 57edcca. No workflow in .github/ runs pytest and the root pyproject declares no testpaths, so nothing in CI changes.
…opologies
deep_research_calc_mathverify{,_partial,_async}.yaml were bring-up scaffolding:
M1 existed to prove the agentic loop "WITHOUT external services (the existing
CalculatorTool + a rule-based MathVerify reward) ... so M2 only swaps in the
search/visit tools and the LLM-judge reward (config-only)". M2 landed and the
scaffold has no remaining purpose.
The catch: all three were the hydra config_name DEFAULTS for the three
answer-graded entrypoints, and search_judge existed in barrier flavour only.
Deleting them would have left train_agentic_partial and train_agentic_async
with no recipe. Repointing both at the barrier recipe is not an option either:
it declares TensorWeightSync, so the async entrypoint would compose and then
fail at runtime, since disaggregated training needs NCCLWeightSync to cross the
slab boundary.
So port M2 to the two missing topologies, grafting each one's topology block off
the calc recipe it replaces and leaving tools/reward/system-prompt identical:
deep_research_search_judge_partial.yaml colocate, TensorWeightSync,
oversample_batch_size 6, partial_rollout, worker_max_concurrency 24
deep_research_search_judge_async.yaml disaggregated, NCCLWeightSync,
train_fraction 0.5, oversample_batch_size 8, mem_fraction_static 0.8
tail_policy: carry is correct for both -- SearchTool and VisitTool subclass the
stateless Tool, not StatefulTool, so a carried Sample holds all resume state
(the same reason the calc recipes carried).
Verified: all six agentic entrypoints compose against their defaults with
--cfg job --resolve, and the resolved configs carry the right weight-sync class,
mem_fraction_static, and prompts_per_rollout == oversample_batch_size per
topology. check-recipe-targets resolves 2135 paths; hooks and bash -n clean.
Neither new recipe has been train-verified -- they are compose-checked ports.
CalculatorTool itself is left in place; no recipe uses it now, but it stays a
supported library tool.
celve
added a commit
to celve/unirl
that referenced
this pull request
Jul 27, 2026
Mirrors the z_image migration (Tencent-Hunyuan#214): generate() becomes the Sample -> Sample endomorphism — read the prompt via sample.conditioning(), take DiffusionSamplingParams (with pinned sigmas) from the frontier gen Part, derive noise groups from frontier.group_ids and x_T via NoiseRecipe.from_sample, then fill the frontier shell with the LatentSegment, decoded images, and replay conditions. User-supplied negatives are deferred (CFG synthesizes the "" negative routed to the DROP system prompt). Adds load_vae for separate-engine recipes (Optional VAE + decode guard).
celve
added a commit
to celve/unirl
that referenced
this pull request
Jul 27, 2026
Mirrors the z_image migration (Tencent-Hunyuan#214): generate() becomes the Sample -> Sample endomorphism — read the prompt via sample.conditioning(), take DiffusionSamplingParams (with pinned sigmas) from the frontier gen Part, derive noise groups from frontier.group_ids and x_T via NoiseRecipe.from_sample, then fill the frontier shell with the LatentSegment, decoded images, and replay conditions. User-supplied negatives are deferred (CFG synthesizes the "" negative routed to the DROP system prompt). Adds load_vae for separate-engine recipes (Optional VAE + decode guard).
celve
added a commit
that referenced
this pull request
Jul 28, 2026
…PO recipe (#219) * feat(boogu_image): vendor Boogu-Image model code verbatim @ 434fb56 Verbatim copy of the 8 model files from boogu-project/Boogu-Image at 434fb56d2a5f6ae6f30cce6a8b56e319e9cd1979 (transformer, attention processors, rope, Lumina2 block, embeddings, components, import/teacache utils). No modifications in this commit; mechanical import rewrites and the SDPA pin follow separately for reviewable diffs. * refactor(boogu_image): vendor mechanical edits — flatten imports, pin SDPA, stub caches - flatten cross-package imports to same-dir relatives - replace triton-RMSNorm / flash-swiglu env conditionals with the upstream default-env branches (torch.nn.RMSNorm, pure-torch swiglu) - pin the four os.getenv("device")-gated attention-processor selections to the SDPA processors; keep Flash2Varlen classes for an optional post-load swap (identical param names, so checkpoints/LoRA stay backend-independent) - route TaylorSeer/TeaCache helper imports to raising stubs (RL rollout and replay must be cache-free); teacache_util vendored verbatim because TeaCacheParams is instantiated unconditionally - add vendor __init__ + VENDOR_COMMIT.txt documenting every edit and the re-vendor recipe * feat(boogu_image): Base T2I model bundle package Adds unirl/models/boogu_image/: config/bundle/conditions/text_embed/ diffusion/vae/pipeline implementing the typed pipeline contract for Boogu-Image-0.1-Base (10.29B vendored DiT + frozen Qwen3-VL-8B mllm encoder + FLUX.1 AutoencoderKL). Key model-specific behavior, each mirroring the reference pipeline: - Qwen3-VL chat-template embed stage with the fixed T2I system prompt; empty/whitespace prompts (the CFG negative "") route to the DROP system prompt; full right-padded sequence + mask, no repacking - predict_noise: t = 1 - sigma in model dtype, positional forward with bare-tensor return, sequential negative branch gated on guidance_scale > 1.0 with plain linear combine, velocity negated to the sigma convention; cfg_range collapses to a per-step effective guidance scale via sampler_kwargs - resolution-independent rotary tables cached per device on the stage - static-v1 schedule pinned as FlowMatchSchedulePolicy.static_only (shift = e^1.15; Boogu's scheduler JSON uses custom field names the base loader would silently ignore) - bundle: mllm lm_head strip, flock-serialized loading, trivial-case meta-init (zero buffers), optional flash2_varlen processor swap * feat(boogu_image): trainside FlowGRPO recipe + einops dep examples/diffusion/boogu_image/boogu_image_trainside.yaml: FlowGRPO + LoRA + PickScore on 8 GPUs at 512^2 (CFG off in Boogu's convention, guidance_scale=1.0; documented 1024^2 fallback overrides since Boogu's native band is 1K-2K). FSDP block_class_names lists the 5 concrete instantiated block classes (exact type-name matching; the DoubleStream block is a separate class, not a subclass). LoRA targets the 12 attention projections incl. the double-stream processor-internal img_*/instruct_* Linears (dot-boundary suffix matching cannot reach them via to_q). master_dtype fp32 per bagel's reward-collapse fix; old_logp_source=replay since forward_batch_size != micro_batch_size. pyproject: add einops>=0.7 (imported by the vendored DiT; pure-python, engine-agnostic). * style(boogu_image): ruff-format glue files; exclude vendored tree from hooks Vendored Boogu code joins bagel/vendor in the pre-commit exclude so re-vendor diffs stay clean against upstream; the three package glue files keep the repo formatting. * refactor(boogu_image): migrate to the Sample-native rollout model Mirrors the z_image migration (#214): generate() becomes the Sample -> Sample endomorphism — read the prompt via sample.conditioning(), take DiffusionSamplingParams (with pinned sigmas) from the frontier gen Part, derive noise groups from frontier.group_ids and x_T via NoiseRecipe.from_sample, then fill the frontier shell with the LatentSegment, decoded images, and replay conditions. User-supplied negatives are deferred (CFG synthesizes the "" negative routed to the DROP system prompt). Adds load_vae for separate-engine recipes (Optional VAE + decode guard).
2 tasks
CjhHa1
added a commit
to zzhuoxin1508/UniRL
that referenced
this pull request
Jul 29, 2026
Brings in the sample-native rollout boundary (Tencent-Hunyuan#214), which removed unirl/types/rollout_req.py and unirl/types/rollout_resp.py along with the RolloutReq / RolloutResp / RolloutTrack triplet. Conflict resolution in unirl/trainer/diffusion.py takes main's version and re-applies this branch's evaluate() seam (sync_weights / sleep_after) plus the _train_fraction field on top of it. unirl/trainer/async_diffusion.py still imports the deleted types at this commit, so it does not import here; the next commit migrates it.
CjhHa1
added a commit
to zzhuoxin1508/UniRL
that referenced
this pull request
Jul 29, 2026
…sync runtime Two things broke this branch against current main, and both are fixed here. The trainer was written against the retired RolloutReq / RolloutResp / RolloutTrack triplet, deleted by the sample-native rollout boundary (Tencent-Hunyuan#214). It is now sample-native: the request is the Sample from _build_request_sample, scoring is reward.score_and_attach(sample) on the self-contained filled Sample instead of the old (req=, track=) pair, groups reassemble with Sample.concat, and the RolloutResp(tracks=...) rebuild and its _track_key bookkeeping are gone. The entry point's stage_config was likewise renamed to task_config. The async buffer / generate seam this branch duplicated from AsyncARTrainer has since been lifted into unirl/rollout/async_runtime.py, the follow-up refactor this PR's description anticipated. _RolloutBuffer, _generate_async, _collect_resp, _is_ready, _launch, _reap_ready and the _next_batch loop are all replaced by AsyncRolloutScheduler + RayGenerationDispatcher, leaving only the diffusion hooks: build a request Sample, score-and-split at reap time, and advantage + FlowGRPO step. Adopting that runtime needs one addition to it, because it launched before it reaped and this path requires the opposite. Reaping pulls the trajectory segment off the rollout slab as an NCCL send issued on the rollout workers, so a generation launched ahead of that send blocks it -- the ~150s/rollout instead of ~8s that reap-before-launch was introduced to fix. Reap-first at max_inflight=1 hands the send idle workers while still launching before the step returns, so the next generation overlaps the caller's train step. The new reap_before_launch flag selects the order and defaults to the existing launch-first behavior, so the AR path is unchanged. Verified: ruff check and format clean, the trainer and entry point import against current main, Hydra compose of the BAGEL async recipe passes, all recipe _target_ paths resolve, both constructor guards fire before any Ray construction, and a fake-dispatcher check confirms reap-first at max_inflight=1 both keeps one generation in flight across every train step and always reaps against idle rollout workers. Not re-run: the GPU reward-curve and localize-timing validation in the PR description.
This was referenced Jul 29, 2026
leviking98z-rgb
added a commit
that referenced
this pull request
Jul 29, 2026
…n RL (#192) * feat(trainer): add AsyncDiffusionTrainer for disaggregated async diffusion RL Diffusion sibling of AsyncARTrainer: subclasses DiffusionTrainer(layout=separate) to reuse the two-slab build + NCCLWeightSync handshake, and overlays the async rollout buffer loop (non-blocking generate, reap-time reward scoring off the train critical path, buffer of scored GRPO groups, train consumes the freshest batch). Knobs: max_inflight (overlap depth), buffer_max_staleness (0=on-policy). Adds unirl/trainer/async_diffusion.py, unirl/train_async_diffusion.py, and examples/diffusion/sd3/sd3_vllmomni_async.yaml. Purely additive. * Overlap async diffusion rollout with training via reap-before-launch segment transfer, add BAGEL async recipe, drop SD3 async recipe * chore(trainer): make train_async_diffusion.py executable to match train_diffusion.py * fix(trainer): point async diffusion entry at BAGEL recipe with stale=2 Default Hydra config still pointed at the dropped SD3 async yaml. Point train_async_diffusion at bagel_vllmomni_async and set buffer_max_staleness=2 (the throughput-optimal knob from the PR validation table). * fix(trainer): keep async diffusion evaluation policy-stable Evaluate the resident rollout policy without syncing or offloading the async engine, while preserving synchronous defaults and forwarding configured eval suites. * docs(trainer): clarify async reward scoring boundary State consistently that generation overlaps training while reap-time reward scoring remains synchronous. * fix(trainer): reject unsupported async diffusion depth Fail before worker construction unless max_inflight is exactly one, preserving the idle-worker window required by reap-time transfer. * fix(trainer): align async diffusion policy metadata Record the train slab fraction and describe the actual remote LoRA sync and bounded policy-lag ratio semantics without changing runtime behavior. * fix(trainer): harden async diffusion result handling * refactor(trainer): async diffusion on the Sample API and the shared async runtime Two things broke this branch against current main, and both are fixed here. The trainer was written against the retired RolloutReq / RolloutResp / RolloutTrack triplet, deleted by the sample-native rollout boundary (#214). It is now sample-native: the request is the Sample from _build_request_sample, scoring is reward.score_and_attach(sample) on the self-contained filled Sample instead of the old (req=, track=) pair, groups reassemble with Sample.concat, and the RolloutResp(tracks=...) rebuild and its _track_key bookkeeping are gone. The entry point's stage_config was likewise renamed to task_config. The async buffer / generate seam this branch duplicated from AsyncARTrainer has since been lifted into unirl/rollout/async_runtime.py, the follow-up refactor this PR's description anticipated. _RolloutBuffer, _generate_async, _collect_resp, _is_ready, _launch, _reap_ready and the _next_batch loop are all replaced by AsyncRolloutScheduler + RayGenerationDispatcher, leaving only the diffusion hooks: build a request Sample, score-and-split at reap time, and advantage + FlowGRPO step. Adopting that runtime needs one addition to it, because it launched before it reaped and this path requires the opposite. Reaping pulls the trajectory segment off the rollout slab as an NCCL send issued on the rollout workers, so a generation launched ahead of that send blocks it -- the ~150s/rollout instead of ~8s that reap-before-launch was introduced to fix. Reap-first at max_inflight=1 hands the send idle workers while still launching before the step returns, so the next generation overlaps the caller's train step. The new reap_before_launch flag selects the order and defaults to the existing launch-first behavior, so the AR path is unchanged. Verified: ruff check and format clean, the trainer and entry point import against current main, Hydra compose of the BAGEL async recipe passes, all recipe _target_ paths resolve, both constructor guards fire before any Ray construction, and a fake-dispatcher check confirms reap-first at max_inflight=1 both keeps one generation in flight across every train step and always reaps against idle rollout workers. Not re-run: the GPU reward-curve and localize-timing validation in the PR description. --------- Co-authored-by: Jianghai <72591262+CjhHa1@users.noreply.github.com> Co-authored-by: Haonan Wang <haonan.wang@u.nus.edu> Co-authored-by: aimicahchen <aimicahchen@tencent.com> Co-authored-by: leviking98z-rgb <leviking98z@gmail.com> Co-authored-by: CjhHa1 <cjh18671720497@outlook.com>
leviking98z-rgb
pushed a commit
to NancyFyong/UniRL
that referenced
this pull request
Jul 29, 2026
…lloutReq The video adapter module still imports unirl.types.rollout_req, which Tencent-Hunyuan#214 removed when Sample/Part replaced the RolloutReq/RolloutResp triplet. The import is unguarded and adapters/__init__ pulls the whole module in for its registration side-effects, so importing the sglang_diffusion engine raises ModuleNotFoundError and every video family goes down with it — WAN 2.1/2.2, HunyuanVideo and Mochi as well as LTX-2. Ltx2T2VAdapter carries two more leftovers behind that import: NoiseRecipe.from_rollout_req, renamed to from_sample, and req.sigmas, a field Sample does not have. Both are ports of what image.py and the VideoAdapter base already do, so the sigma schedule now comes from sample.frontier_gen_part(DiffusionSamplingParams).sampling_params.sigmas.
leviking98z-rgb
added a commit
that referenced
this pull request
Jul 29, 2026
…lloutReq (#272) The video adapter module still imports unirl.types.rollout_req, which #214 removed when Sample/Part replaced the RolloutReq/RolloutResp triplet. The import is unguarded and adapters/__init__ pulls the whole module in for its registration side-effects, so importing the sglang_diffusion engine raises ModuleNotFoundError and every video family goes down with it — WAN 2.1/2.2, HunyuanVideo and Mochi as well as LTX-2. Ltx2T2VAdapter carries two more leftovers behind that import: NoiseRecipe.from_rollout_req, renamed to from_sample, and req.sigmas, a field Sample does not have. Both are ports of what image.py and the VideoAdapter base already do, so the sigma schedule now comes from sample.frontier_gen_part(DiffusionSamplingParams).sampling_params.sigmas. Co-authored-by: LeviKing <leviking98z@gmail.com>
haonan3
added a commit
to YSunLIN/UniRL
that referenced
this pull request
Jul 30, 2026
…oles DSL - Rebase onto current main (Tencent-Hunyuan#214): RolloutReq/RolloutInputs are gone; the actor role now consumes Texts/Images primitives plus per-sample metadata records straight from the data-source Sample. - Drop recipes/common/ (role-list orchestration): REFLTrainer subclasses BaseTrainer directly and wires actor + reward with placement()+remote_hydra, mirroring RewardBackpropTrainer (the SD3 image-ReFL driver). ReflActorRole mirrors ReFLPolicy's family-agnostic contract (pipeline_target + model_config + from_config; FSDPBackend composed in initialize()). - Re-root both configs from the roles: list to the repo-wide flat schema (actor:/reward:/data_source:/sampling:/logging:), the same shape as examples/diffusion/refl_sd3.yaml. - KL correctness across DP shards: diffuse_with_grad now returns per-sample [B] KL (concat field) instead of a per-shard scalar shared field, so DP_SCATTER merge/re-shard round-trips each shard's own KL. Previously the driver collapsed all shards to one value (loss/logging skew; gradient flow was unaffected because dKL/dkl is the constant kl_weight). - I2V condition assembly moves out of the role into Wan21/Wan22ReflPipeline.build_refl_conditions (mirrors each mainline pipeline's generate); negative prompts ride sampler_kwargs. - Seed scheme now matches ReFLPolicy (base + 1000*rollout_id + dp_rank); previously every rollout redrew the same init noise. - reward/service.py: restore the List typing import lost in the merge.
5 tasks
2 tasks
CjhHa1
added a commit
to HaitaoWuTJU/UniRL
that referenced
this pull request
Jul 31, 2026
The package did not import: Tencent-Hunyuan#214 replaced RolloutReq / RolloutResp with Sample / Part, and the main merge into this branch did not conflict because pipeline.py exists only on the feature side, so it kept importing unirl.types.rollout_req. Every validation number in the PR description was produced before that merge. generate is now Sample -> Sample. Task selection reads parts[0].control["task"] and otherwise infers from Sample.has_image_input(). i2t pulls its turns from vision_conditioning() and t2i from turns(); both fail loud on a multi-turn trajectory, which the single-user-turn Janus chat template cannot encode. Sampling params come off the frontier gen shell rather than a request-level dict, and t2i requires JanusProImageARSamplingParams because the image grid, CFG weight, and token count all ride on it. Results go back through Sample.with_filled_frontier, so the input chain and reward_compute_s survive. The stages, conditions, bundle, config, chat_template, and image_prompt were already API-agnostic and are untouched.
2 tasks
3 tasks
haonan3
added a commit
to KemingWu/UniRL
that referenced
this pull request
Aug 2, 2026
…rework) Adopts the core-primitive DiffusionOPD implementation in place of the original implementation on this branch; the tree is taken entirely from the rework side (-s ours), the original commits stay in history. What changed relative to the original approach: - teachers are backend-owned frozen LoRA adapters (backend.lora_cfg.frozen_adapters, injected pre-FSDP-wrap, so FSDP shards them and LoRA checkpoints stay save/load-symmetric) - teacher selection is data-driven via metadata["domain"] stamped by MultiDomainRLDataSource (replaces the data-source/algorithm round-robin counter pair, which desynchronized on checkpoint resume and on the branch's post-Tencent-Hunyuan#214 rebase broke entirely: unirl.types.prompts was removed upstream) - eval covers every domain in one pass (budget split across domains) - local mmdet GenEval scorer and the 2-node launcher snapshot are dropped; classical GenEval stays in unirl-reward-service (configs/geneval_service.yaml)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This draft proposes a Sample-native agentic RL integration for UniRL.
Sample -> Sample/Sample -> list[Sample]contract and migrate existing model, rollout, reward, and training paths ontoSample/Part.The draft is intentionally opened for architecture and sequencing review before the remaining final-HEAD GPU matrix and patch-size cleanup are complete.
Test Plan
Completed on the current draft tip (
159627bd):SKIP=no-commit-to-branch pre-commit run --all-files --show-diff-on-failurepython3 -m compileall -q unirlpython3 scripts/check_recipe_targets.py— 1,965 Hydra_target_paths resolvedOmegaConf.resolvefor all six barrier/partial/async ALFWorld and calculator/deep-research recipesReal-GPU evidence from integration ancestors:
b4ed9375).9cd7f3ab).ce13a48d).159627bd) on 8× H20 with a local PickScore smoke override: 8 source-image prompts × 2 edits, the three-Part text → source image → generated image lineage and typedimage_latent/image_latent_idsreplay conditions passed on every rank, and one optimizer step completed withtrain/has_backward=1and nonzero grad norm (W&B).diffusion/bagel/bagel_editreward) completed on 8× H20 with EditReward isolated on a second pod: the default step-0 deterministic evaluation ran, followed by 3 full rollouts at the default 8 prompts × 8 edits (192 training edits total), 14 denoising steps, and 2 optimizer updates per rollout (6 total). All 32 recipe scoring batches, plus the separate two-turn preflight, returned HTTP 200; rollout rewards were-0.7940,-0.9308, and-0.7629; and eight 1024×512 source|edited previews synced successfully (W&B). This validation used two working-tree fixes not yet committed to the PR: hydrating the source-imageTensorRefbefore media-preview indexing and treating EditReward's blankoutput_diras its/tmpinference default.Pending before ready-for-review:
top_kbehavior.Compatibility / Risk
RolloutReq,RolloutResp, andRolloutTrackare removed without compatibility aliases. Seeunirl/types/README.mdfor field/helper mappings.Reviewer Notes
Suggested review order:
unirl/types/Sample/Part lineage and migration contract.unirl/rollout/loop/and tools.Overlap/sequence needs maintainer agreement:
RolloutResparchitecture; this draft proposes the Sample/Part direction instead.Temporary integration-only smoke scripts and the top-level test tree used during development were intentionally removed from the upstream patch; the retained validation evidence is summarized above. Before ready-for-review, we will complete the final-tip GPU matrix and perform a subsystem-by-subsystem human review. This integration was developed with substantial AI assistance (Codex); the human submitter will review and defend the final diff before the draft is promoted.
Checklist
Current-head maintainer validation (28a27e4)
main(fce4f4b) and resolved the Sample/Part migration conflicts; GitHub now reports the PR mergeable./tmp/editreward_outputwhen the vendored config leavesoutput_dirblank.maincontains notests/**changes.pytest tests -q(25 passed); 32 deleted historical PR regression cases executed in memory (all passed, no files restored);compileall; full pre-commit; all seven new agentic recipes and four modified diffusion recipes composed and resolved.pytestalso runs the same 25 tests successfully, then mis-collects six existing command-line functions inunirl-reward-service/scripts/test_videoalign.pybecause they require CLI arguments; that script is unchanged relative tomain.This removes the textual merge conflict and known uncommitted BAGEL fixes, but does not make the patch low risk. Final-tip agentic/AR/BAGEL GPU reruns, dependency-matrix evidence, patch-size cleanup, and subsystem human review remain required before merge.