fix(rollout): guard agentic sync and deterministic diffusion eval

unirl/rollout/README.md

@@ -96,7 +96,9 @@ implements its weight-receive method and a matching `sync:` handler in
   `AsyncBatchRolloutEngine.quiesce()` drains every in-flight generation; a
   weight + KV update corrupts one mid-flight. The agentic quiesce is a
   turn-boundary `abort` + final poll, folded into
-  `AsyncAgenticRolloutEngine.quiesce()`. Reap-vs-launch ordering is trainer
+  `AsyncAgenticRolloutEngine.quiesce()`; its `sync_weights()` rejects a live
+  drive, then pairs the weight push with the version bump and logs the sync.
+  Reap-vs-launch ordering is trainer
   statement order (diffusion polls before topping up; see its `_next_step`).
 - **Reward/advantage methods are not engine code** — `Part.compute_advantages` and
   `Sample.propagate_rewards` are called by the trainer after scoring. An engine

unirl/rollout/engine/asynchronous.py

@@ -11,7 +11,7 @@
 - :class:`VersionedBuffer` — payload-agnostic freshness/staleness buffer.
 - :class:`InflightPool` — non-blocking pool of distributed ``generate`` calls.
 
-Engines, sharing one consumer surface (``poll`` / ``drain_freshest`` /
+Engines share one consumer surface (``poll`` / ``drain_freshest`` /
 ``pop_evicted`` / ``quiesce`` + engine-owned ``weight_version``):
 
 - :class:`AsyncBatchRolloutEngine` — batch granularity over a single-turn
@@ -336,8 +336,19 @@ def __init__(self, rollout: Any, *, group_size: int, start_gen_id: int = 0) -> N
     def weight_version(self) -> int:
         return self._weight_version
 
-    def bump_weight_version(self) -> int:
+    def sync_weights(self, weight_sync: Any) -> int:
+        """Push train weights via *weight_sync* and advance the version ledger.
+
+        The only sanctioned weight-push path — pairing the push with the bump
+        is what keeps the ledger truthful. Raises while a drive is active (a
+        weight push must be decode-idle); a joined ``finalize_if_drained`` or
+        ``quiesce`` ends the drive.
+        """
+        if self._drive_live:
+            raise RuntimeError("sync_weights with a drive active; finalize or quiesce() first")
+        weight_sync.sync()
         self._weight_version += 1
+        logger.info("sync_weights: pushed train weights; weight_version -> %d", self._weight_version)
         return self._weight_version
 
     def submit(self, tasks: List["Sample"]) -> None:
@@ -380,7 +391,7 @@ def pop_evicted(self) -> List[List["Sample"]]:
     def quiesce(self) -> List["Sample"]:
         """Turn-boundary stop: abort, then one final poll for trajectories that
         completed DURING the quiesce (before the next ``submit`` resets worker
-        buffers). Call before ``bump_weight_version`` so those groups carry the
+        buffers). Call before ``sync_weights`` so those groups carry the
         version they completed under."""
         carried = self._rollout.abort()[0]
         self.poll()

unirl/trainer/agentic_async.py

@@ -21,8 +21,8 @@
   ``buffer_max_staleness``), reward + GRPO advantage + one optimizer step (reusing
   :class:`AgenticTrainer`'s helpers), then **quiesce + sync**: ``abort`` the in-flight
   tail at a turn boundary, apply the configured ``tail_policy`` (carry only when the
-  environment can resume from the ``Sample``; otherwise drop), ``weight_sync.sync()``,
-  bump the version.
+  environment can resume from the ``Sample``; otherwise drop), then
+  ``engine.sync_weights`` (one call: push + version bump).
 
 ONE single-threaded loop (the ``AsyncARTrainer`` shape): with disjoint slabs the
 rollout slab keeps generating in the background (the engine's per-worker drain) while
@@ -356,7 +356,7 @@ def train(
         self._engine = AsyncAgenticRolloutEngine(self.rollout, group_size=self._n, start_gen_id=start_rollout)
 
         if start_rollout < num_rollouts and start_rollout and self.weight_sync is not None:
-            self.weight_sync.sync()
+            self._engine.sync_weights(self.weight_sync)  # push restored weights into the fresh engine
         if start_rollout < num_rollouts:
             self._submit_drive(carried=[], rollout_id=start_rollout)
 
@@ -387,8 +387,7 @@ def train(
                             save_mode=save_mode,
                         )
                     if need_sync:
-                        self.weight_sync.sync()
-                        self._engine.bump_weight_version()
+                        self._engine.sync_weights(self.weight_sync)
                     if step < num_rollouts:
                         self._submit_drive(carried=carried, rollout_id=step)
         finally:

unirl/trainer/agentic_partial.py

@@ -155,8 +155,7 @@ def _collect_until(self, batch_size: int, rollout_id: int, stale: int) -> List[L
     def _drive_partial(self, rollout_id: int, sync_weights: bool, stale: int) -> List[List[Sample]]:
         self.rollout.wake_up()
         if sync_weights and self.weight_sync is not None:
-            self.weight_sync.sync()
-            self._engine.bump_weight_version()
+            self._engine.sync_weights(self.weight_sync)
         tasks = self._build_tasks(self._carried, rollout_id)
         self._carried = []
         self._engine.submit(tasks)

unirl/trainer/diffusion.py

@@ -388,8 +388,9 @@ def evaluate(
 
         Mirrors :meth:`train_step`'s rollout+reward path but skips advantage/backward.
         Generates at the deterministic best-quality setting (``cfg_text_scale=
-        eval_cfg_text_scale``, ``eta=eval_eta``; ``eval_samples_per_prompt`` x_T per
-        prompt) and scores. The training reward plus every shared-set
+        eval_cfg_text_scale``, ``eta=eval_eta`` — at ``eval_eta=0`` the SDE gate
+        is also cleared, so the request is pure ODE; ``eval_samples_per_prompt``
+        x_T per prompt) and scores. The training reward plus every shared-set
         ``eval_rewards`` suite scores the SAME generated images over the default
         eval set (``run.eval_data_path``, ``eval_num_prompts`` prompts); each
         own-set suite then gets its own generation pass over its own prompts.
@@ -407,6 +408,11 @@ def evaluate(
             samples_per_prompt=self.eval_samples_per_prompt,
             eta=self.eval_eta,
         )
+        if self.eval_eta <= 0.0:
+            # Deterministic eval must also clear the SDE gate: eta=0 with gated
+            # steps is a contradictory request — the central kernel degrades such
+            # steps to ODE, but worker-resident schedulers (BAGEL) refuse the pair.
+            replace_kwargs.update(sde_indices=[], scheduler=None)
         if "cfg_text_scale" in {f.name for f in dataclasses.fields(base_diffusion)}:
             replace_kwargs["cfg_text_scale"] = self.eval_cfg_text_scale
         else: