feat(rollout): let reward actors read the generated output from the output store

docs/advanced/streaming-reward.md

@@ -21,9 +21,8 @@ flowchart LR
 
 
 The rollout pipeline decomposes into `generate` / `deserialize` / `reward` components. With `--use-output-store`, the
-output is read back from the store and deleted there between `deserialize` and `reward`, still inside the generation
-slot. Because every microgroup is its own `asyncio` task, the engine starts the next request while previous responses
-are still being unpacked and scored.
+reward actors read the output from the store, and it is deleted there once scored. Because every microgroup is its
+own `asyncio` task, the engine starts the next request while previous responses are still being unpacked and scored.
 Nothing waits for the full batch.
 
 ## 2. Deserialization: msgpack + parser-actor pool

docs/user-guide/customization.md

@@ -57,6 +57,10 @@ def generate_rollout(
 Legacy list/dict returns are still wrapped into the corresponding output dataclass by
 `miles.rollout.base_types.call_rollout_fn`.
 
+With `--use-output-store`, a rollout function that does not go through `generate_rollout_async` sets
+`GenerateState(args).media_outputs_left` to the number of outputs to keep for its media logging at the start of each
+rollout, and calls `release_generated_outputs` after a reward it runs itself.
+
 ### `--eval-function-path`
 
 Same shape as the rollout function. Defaults to `--rollout-function-path` when
@@ -92,10 +96,8 @@ async def custom_generate(
 
 `evaluation=` is optional — if your signature omits it, the caller skips that kwarg.
 
-With `--use-output-store`, `generate_microgroup` leaves `sample.generated_output` as a `StoredTensor`.
-`generate_and_rm_microgroup` reads it into the tensor after your function returns and before the reward. Code that
-reads the output earlier, such as a custom generate function that refines its own first pass, calls
-`await load_generated_outputs(GenerateState(args), microgroup)` first.
+With `--use-output-store`, `generate_microgroup` leaves `sample.generated_output` as a `StoredTensor`: the output
+stays in the store for the reward actors, and a custom generate function cannot read it.
 
 ### `--diffusion-step-strategy-path`
 
@@ -134,6 +136,13 @@ Built-in scorers (`--rm-type pickscore` / `hps` / `ocr` / `custom_api` / `openai
 [Rewards](rewards.md). `--custom-rm-path` replaces that dispatch entirely; `--custom-reward-post-process-path`
 replaces the normalization that runs after scoring.
 
+With `--use-output-store`, the rollout process does not read a training rollout's output. There,
+`sample.generated_output` is a `StoredTensor`, and the actors of a reward pool read it when they score it. A
+`--custom-rm-path` function therefore scores through a pool, a built-in `*_rm` or its own `AsyncRewardActorPool`, and
+does not read the tensor itself. Once a sample is scored its output is `None`; the first `--wandb-log-num-images`
+scored samples of a rollout that logs media keep it as a tensor, so a rollout whose first groups a filter drops logs
+fewer. Every hook that runs after scoring sees that.
+
 ### `--custom-rm-path`
 
 ```python
@@ -183,6 +192,9 @@ async def api_rm(args, samples: list[Sample], **kwargs) -> list[float]:
 --rm-url http://localhost:8000/score
 ```
 
+With `--use-output-store` this function cannot read the output. Send it through `--custom-api-rm-config`, whose
+`ApiRewardActor` reads it in the actor, or through your own `AsyncRewardActorPool(args=args, ...)`.
+
 ### `--custom-reward-post-process-path`
 
 Replace GRPO advantage normalization in

miles/dashboard/events.py

@@ -33,16 +33,13 @@ class TrajectoryEventKind(StrEnum):
     DESER_END = "deser_end"
     REWARD_START = "reward_start"
     REWARD_END = "reward_end"
-    STORE_READ_START = "store_read_start"
-    STORE_READ_END = "store_read_end"
 
 
 # stage name (StageTimer) -> (start kind, end kind)
 STAGE_KINDS = {
     "generate": (TrajectoryEventKind.GEN_START, TrajectoryEventKind.GEN_END),
     "deserialize": (TrajectoryEventKind.DESER_START, TrajectoryEventKind.DESER_END),
     "reward": (TrajectoryEventKind.REWARD_START, TrajectoryEventKind.REWARD_END),
-    "store_read": (TrajectoryEventKind.STORE_READ_START, TrajectoryEventKind.STORE_READ_END),
 }
 
 # kind -> (base span, is_start), for the reader to fold start/end back into spans
@@ -53,8 +50,6 @@ class TrajectoryEventKind(StrEnum):
     "deser_end": ("deser", False),
     "reward_start": ("reward", True),
     "reward_end": ("reward", False),
-    "store_read_start": ("store_read", True),
-    "store_read_end": ("store_read", False),
 }
 
 

miles/ray/rollout.py

@@ -449,12 +449,15 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl
         reward_stats["rollout/step"] = compute_rollout_step(self.args, self.rollout_id)
         tracking_utils.log(self.args, reward_stats, step_key="rollout/step")
 
-        max_images = self.args.wandb_log_num_images
-        interval = self.args.wandb_log_image_interval
-        if max_images > 0 and self.rollout_id % interval == 0:
+        if tracking_utils.logs_rollout_media(self.args, self.rollout_id):
+            # with the output store on, only the samples the rollout kept for media hold a tensor
             self._log_images_async(
-                {"rollout_media/sample_images": samples},
-                max_images=max_images,
+                {
+                    "rollout_media/sample_images": [
+                        sample for sample in samples if isinstance(sample.generated_output, torch.Tensor)
+                    ]
+                },
+                max_images=self.args.wandb_log_num_images,
                 step_key="rollout/step",
                 step_value=compute_rollout_step(self.args, self.rollout_id),
                 reward_key=self.args.reward_key,

miles/rollout/rm_hub/api.py

@@ -66,6 +66,7 @@ def __init__(self, args) -> None:
         if not isinstance(actor_cls, type) or not issubclass(actor_cls, self.actor_base_cls):
             raise TypeError(f"API reward actor_class must be an {self.actor_base_cls.__name__} subclass")
         super().__init__(
+            args=args,
             actor_cls=actor_cls,
             actor_kwargs=config.actor_kwargs,
             num_workers=1,

miles/rollout/rm_hub/core.py

@@ -9,6 +9,8 @@
 from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
 
 from miles.ray.utils import COLOCATED_REWARD_GPU
+from miles.utils.diffusion_rollout_response import read_generated_outputs
+from miles.utils.output_store import RegisteredTensorReader, StoredTensor, connect_output_store
 
 logger = logging.getLogger(__name__)
 
@@ -79,10 +81,27 @@ def record_reward_queue_depth(samples, name: str, max_queue_depth: int) -> None:
         sample.reward_max_queue_depth = depths
 
 
+class _StoredOutputRewardActor:
+    """A reward actor that reads the stored outputs it is given from the output store, then scores them."""
+
+    def __init__(self, *, scorer_cls, scorer_kwargs: dict, args) -> None:
+        self._scorer = scorer_cls(**scorer_kwargs)
+        self._reader = RegisteredTensorReader(connect_output_store(args))
+
+    def score_batch(self, outputs: list, prompts: list[str]) -> list[float]:
+        # eval outputs arrive inline
+        stored_outputs = [output for output in outputs if isinstance(output, StoredTensor)]
+        if stored_outputs:
+            read_outputs = iter(read_generated_outputs(self._reader, stored_outputs))
+            outputs = [next(read_outputs) if isinstance(output, StoredTensor) else output for output in outputs]
+        return self._scorer.score_batch(outputs, prompts)
+
+
 class AsyncRewardActorPool:
     """Round-robin pool for Ray reward actors exposing ``score_batch(outputs, prompts)``.
 
-    Actors receive the float ``generated_output`` tensors and quantise to uint8 themselves.
+    Actors receive the float ``generated_output`` tensors and quantise to uint8 themselves. With the output store
+    on they are given stored outputs and read them first.
     Colocated pools take one slot per rollout bundle from the manager's
     ``slots``; standalone pools are default-scheduled at ``num_gpus_per_worker``, which
     only lands on GPUs outside every placement group.
@@ -91,6 +110,7 @@ class AsyncRewardActorPool:
     def __init__(
         self,
         *,
+        args,
         actor_cls,
         actor_kwargs: dict,
         num_workers: int,
@@ -113,6 +133,9 @@ def __init__(
             num_gpus_per_worker = COLOCATED_REWARD_GPU
         else:
             strategies = ["DEFAULT"] * num_workers
+        if args.use_output_store:
+            actor_kwargs = {"scorer_cls": actor_cls, "scorer_kwargs": actor_kwargs, "args": args}
+            actor_cls = _StoredOutputRewardActor
 
         self._actors = [
             ray.remote(actor_cls)

miles/rollout/rm_hub/hps.py

@@ -128,6 +128,7 @@ class AsyncHPSPool(AsyncRewardActorPool, metaclass=SingletonMeta):
 
     def __init__(self, args, placement_group=None, slots=None) -> None:
         super().__init__(
+            args=args,
             actor_cls=HPSRewardActor,
             actor_kwargs={
                 "hps_version": args.hps_version,

miles/rollout/rm_hub/hpsv3.py

@@ -233,6 +233,7 @@ class AsyncHPSv3Pool(AsyncRewardActorPool, metaclass=SingletonMeta):
 
     def __init__(self, args, placement_group=None, slots=None) -> None:
         super().__init__(
+            args=args,
             actor_cls=HPSv3RewardActor,
             actor_kwargs={"checkpoint_path": args.hpsv3_checkpoint_path},
             num_workers=args.hpsv3_num_workers,

miles/rollout/rm_hub/ocr.py

@@ -99,6 +99,7 @@ class AsyncOcrPool(AsyncRewardActorPool, metaclass=SingletonMeta):
 
     def __init__(self, args) -> None:
         super().__init__(
+            args=args,
             actor_cls=OcrRewardActor,
             actor_kwargs={"use_gpu": False},
             num_workers=args.ocr_num_workers,

miles/rollout/rm_hub/pickscore.py

@@ -7,6 +7,7 @@
 from PIL import Image
 
 from miles.utils.misc import SingletonMeta
+from miles.utils.output_store import StoredTensor
 from miles.utils.processing_utils import generated_output_to_rgb_hwc_uint8_frames, sample_frame_indices
 from miles.utils.types import Sample
 
@@ -74,8 +75,10 @@ def __init__(
         processor_path: str,
         model_path: str,
         frames_per_forward: int,
+        num_frames: int | None,
     ) -> None:
         self.frames_per_forward = frames_per_forward
+        self.num_frames = num_frames
         gpu_ids = ray.get_gpu_ids()
         use_cuda = bool(gpu_ids) and torch.cuda.is_available()
         if use_cuda:
@@ -88,10 +91,10 @@ def __init__(
         )
 
     def score_batch(self, outputs: list[torch.Tensor], prompts: list[str]) -> list[float]:
-        # a video sample is scored on every frame it arrives with and gets their mean
+        # a video is scored on --pickscore-num-frames of its frames, averaged; already picked frames pass unchanged
         images, frame_counts = [], []
         for output in outputs:
-            frames = generated_output_to_rgb_hwc_uint8_frames(output, None)
+            frames = generated_output_to_rgb_hwc_uint8_frames(output, self.num_frames)
             images.extend(Image.fromarray(frame) for frame in frames)
             frame_counts.append(len(frames))
         flat_prompts = [p for p, n in zip(prompts, frame_counts, strict=True) for _ in range(n)]
@@ -112,11 +115,13 @@ class AsyncPickScorePool(AsyncRewardActorPool, metaclass=SingletonMeta):
 
     def __init__(self, args, placement_group=None, slots=None) -> None:
         super().__init__(
+            args=args,
             actor_cls=PickScoreRewardActor,
             actor_kwargs={
                 "processor_path": args.pickscore_processor_path,
                 "model_path": args.pickscore_model_path,
                 "frames_per_forward": args.pickscore_batch_size,
+                "num_frames": args.pickscore_num_frames,
             },
             num_workers=args.pickscore_num_workers,
             batch_size=args.pickscore_batch_size,
@@ -130,10 +135,15 @@ def __init__(self, args, placement_group=None, slots=None) -> None:
 
 async def pickscore_rm(args, samples: Sequence[Sample]) -> list[float]:
     pool = AsyncPickScorePool(args)
-    # pick --pickscore-num-frames here so only the scored frames cross the object store
+    # pick --pickscore-num-frames here so only the scored frames cross the object store; a stored output is not
+    # here to pick from, so its actor picks them
     outputs = [
-        s.generated_output[:, sample_frame_indices(s.generated_output.shape[1], args.pickscore_num_frames)]
-        for s in samples
+        (
+            output
+            if isinstance(output, StoredTensor)
+            else output[:, sample_frame_indices(output.shape[1], args.pickscore_num_frames)]
+        )
+        for output in (s.generated_output for s in samples)
     ]
     scores, max_queue_depth = await pool.score(outputs, [s.prompt for s in samples])
     record_reward_queue_depth(samples, "pickscore", max_queue_depth)

miles/rollout/sglang_diffusion_rollout.py

@@ -22,7 +22,8 @@
 from miles.utils.eval_config import EvalDatasetConfig
 from miles.utils.http_utils import post
 from miles.utils.misc import SingletonMeta, load_function
-from miles.utils.output_store import StoredTensor, connect_output_store, record_dropped_samples
+from miles.utils.output_store import RegisteredTensorReader, StoredTensor, connect_output_store, record_dropped_samples
+from miles.utils.tracking_utils import logs_rollout_media
 from miles.utils.types import Sample
 
 from .rm_hub import batched_async_rm
@@ -132,6 +133,9 @@ def __init__(self, args: Namespace) -> None:
         self._parser_rr = 0
         self.parser_inflight = [0] * args.rollout_parser_num_workers
         self.output_store = connect_output_store(args) if args.use_output_store else None
+        self.output_reader = RegisteredTensorReader(self.output_store) if args.use_output_store else None
+        # outputs of the running rollout still to keep for its media logging
+        self.media_outputs_left = 0
 
         self.reset()
 
@@ -249,33 +253,44 @@ async def generate_microgroup(
     return microgroup
 
 
-async def load_generated_outputs(state: GenerateState, samples: list[Sample]) -> None:
-    """Read the samples' stored outputs into tensors and delete them from the store.
+async def release_generated_outputs(state: GenerateState, samples: list[Sample]) -> None:
+    """Delete the stored outputs of scored samples from the store, first reading the few media logging keeps.
 
-    The store then holds an output only from the engine's put to this read. The rollout's release deletes whatever a
-    failed delete leaves.
+    The samples' outputs become ``None``, except the kept ones, which become tensors. The rollout's release deletes
+    whatever a failed delete leaves.
     """
-    stored = [sample for sample in samples if isinstance(sample.generated_output, StoredTensor)]
-    if not stored:
+    stored_samples = [sample for sample in samples if isinstance(sample.generated_output, StoredTensor)]
+    if not stored_samples:
         return
-    outputs = [sample.generated_output for sample in stored]
-    refs = [
-        msgspec.structs.replace(
-            output.ref, fields=tuple(spec for spec in output.ref.fields if spec.name == output.field)
-        )
-        for output in outputs
-    ]
-    st = hooks.StageTimer()
-    with st.stage("store_read"):
-        tensors = await asyncio.to_thread(read_generated_outputs, state.output_store, outputs)
+    # taken before the read is awaited, so concurrent microgroups keep at most --wandb-log-num-images in all
+    media_samples = stored_samples[: state.media_outputs_left]
+    state.media_outputs_left -= len(media_samples)
+    media_tensors = []
+    if media_samples:
+        media_outputs = [sample.generated_output for sample in media_samples]
         try:
-            await asyncio.to_thread(state.output_store.delete_many, refs)
+            media_tensors = await asyncio.to_thread(read_generated_outputs, state.output_reader, media_outputs)
         except Exception:
             logger.warning(
-                "Failed to delete %d read outputs; the rollout's release deletes them", len(refs), exc_info=True
+                "Failed to read %d outputs kept for media; they are not logged", len(media_samples), exc_info=True
             )
-    st.attach(stored)
-    for sample, tensor in zip(stored, tensors, strict=True):
+            media_samples = []
+    output_refs = [
+        msgspec.structs.replace(
+            stored_output.ref,
+            fields=tuple(spec for spec in stored_output.ref.fields if spec.name == stored_output.field),
+        )
+        for stored_output in (sample.generated_output for sample in stored_samples)
+    ]
+    try:
+        await asyncio.to_thread(state.output_store.delete_many, output_refs)
+    except Exception:
+        logger.warning(
+            "Failed to delete %d scored outputs; the rollout's release deletes them", len(output_refs), exc_info=True
+        )
+    for sample in stored_samples:
+        sample.generated_output = None
+    for sample, tensor in zip(media_samples, media_tensors, strict=True):
         sample.generated_output = tensor
 
 
@@ -299,8 +314,6 @@ async def generate_and_rm_microgroup(
                     microgroup = await custom_generate_func(args, microgroup, sampling_params)
             else:
                 microgroup = await generate_microgroup(args, microgroup, sampling_params, evaluation=evaluation)
-        # inside the semaphore, so the store holds at most one microgroup's outputs per generation slot
-        await load_generated_outputs(state, microgroup)
 
     # for the rm that need the whole group, we will not do the rm here
     if args.group_rm:
@@ -314,13 +327,14 @@ async def generate_and_rm_microgroup(
     for sample, reward in zip(microgroup, rewards, strict=True):
         sample.reward = reward
         hooks.record_trajectory(sample)
+    await release_generated_outputs(state, microgroup)
     return microgroup
 
 
 async def generate_and_rm_group(
     args: Namespace, group: list[Sample], sampling_params: dict[str, Any], evaluation: bool = False
 ) -> list[Sample]:
-    GenerateState(args)
+    state = GenerateState(args)
 
     # N-spaced base so sgl-d's seed→[seed+0..seed+N-1] expansion stays disjoint
     # per (rollout, prompt-group); group_index is monotonic across the run.
@@ -347,6 +361,7 @@ async def generate_and_rm_group(
         rewards = await batched_async_rm(args, group)
         for sample, reward in zip(group, rewards, strict=False):
             sample.reward = reward
+        await release_generated_outputs(state, group)
 
     return group
 
@@ -369,6 +384,7 @@ async def generate_rollout_async(
     assert args.rollout_global_dataset
 
     state = GenerateState(args)
+    state.media_outputs_left = args.wandb_log_num_images if logs_rollout_media(args, rollout_id) else 0
 
     # instantiate data filters
     dynamic_filter = (