[RFC] Add Sample-level Logging API

[DNXie]

This issue propose adding sample-level logging to complement existing scalar metrics, so users can directly inspect concrete prompts, responses, targets, and rewards produced during RL training.

This RFC is structured in two stages.

• Stage 1 introduces the core API for sample-level logging, enabling users to record and inspect concrete prompts, responses, targets, and rewards alongside existing scalar metrics. The focus here is to provide a minimal, consistent mechanism without adding configuration complexity.
• Stage 2 looks ahead to user-defined filters, allowing researchers to control what subset of samples are logged (e.g., by reward thresholds, frequency, or custom predicates). This keeps Stage 1 lightweight while outlining a clear path for more flexible logging in the future.

---

Stage 1: Core API for Sample-level Logging

Motivation

As discussed in #187, we’re currently using reward/loss as our main proxy for RL correctness. It makes it hard to tell if the RL loop is truly learning versus simply reward hacking. Numeric metrics alone don’t show the actual examples that led to a particular reward signal.

By logging structured sample-level information (e.g., prompt, response, target, reward, advantage), users can directly examine the concrete cases that produced specific reward signals. This makes it easier to validate that rewards align with the intended task rather than artifacts, and enables qualitative checks such as:

• Debugging zero-reward responses
• Spotting distribution drift in generated responses.
• Inspecting sample-level statistics (e.g., rewards).

Benefits: This enables qualitative inspection of responses in WandB tables with minimal overhead, since it only logs structured dicts at rollout time. It also stays consistent with the existing metrics system by reusing the record_metric API.

Scope

• Only enable the API for sample-level logging, not token-level.
• Loss is not included: see reasons below.

Design

What to Log

episode_id: int
policy_version: int
prompt: str
response: str
request_len: int
response_len: int
target: str
reward: float
ref_logprobs: float
advantage: float

Alternative Considered

• Log per-sample loss (loss, policy_loss, kl_loss):
  Not in scope for this proposal. Supporting this would require structural changes to the loss function, forward_backward, and potentially train_step. Currently the loss function only returns batch-level loss; to enable per-sample logging, it would need to return a list of per-sample losses, which would then be propagated and logged inside train_step.

Form of Log

• ConsoleBackend → Pretty-print dicts via pprint.
  Example:

  === SAMPLE LOGS STEP 42 ===
  {
    'id': 1,
    'prompt': '2+2?',
    'response': '<answer>5</answer>',
    'target': '4',
    'reward': 0.0,
    'advantage': -1.2
    ...
  }
  

• WandbBackend → wandb.Table
  Columns: [id, prompt, response, target, reward, advantage, ...]
  Rows: one per logged sample, browsable in the WandB UI.

Where to Insert the Log

• Preferred: continuous_rollouts (has full context including advantage and ref_logprobs).
• Alternative: RewardActor.evaluate_response (if we don't want advantage and ref_logprobs).

How to Log

We reuse the existing record_metric API with a new Reduce.SAMPLE:

record_metric(
    "reward/samples",
    {"episode_id": ep.episode_id, "prompt": prompt, ...},
    Reduce.SAMPLE,
)

New Accumulator: SampleAccumulator

class SampleAccumulator(MetricAccumulator):
    def __init__(self, reduction: Reduce):
        super().__init__(reduction)
        self.samples = []

    def append(self, value: dict) -> None:
        assert isinstance(value, dict)
        self.samples.append(value)

    def get_value(self) -> list[dict]:
        return self.samples

    def get_state(self) -> Dict[str, Any]:
        return {
            "reduction_type": self.reduction_type.value,
            "samples": self.samples,
        }

    @classmethod
    def get_reduced_value_from_states(cls, states: List[Dict[str, Any]]) -> list[dict]:
        # Each state looks like {"reduction_type": "sample", "samples": [ ... ]}
        merged = []
        for s in states:
            merged.extend(s.get("samples", []))
        return merged

    def reset(self) -> None:
        self.samples = []

Flush Behavior

async def flush(self, step: int, return_state: bool = False):
    ...
    metrics = {}
    samples = {}   # 👈 new

    for key, state in states.items():
        reduction = state["reduction_type"]
        acc_class = Reduce(reduction).accumulator_class

        if reduction == Reduce.SAMPLE.value:
            samples[key] = acc_class.get_reduced_value_from_states([state])   # 👈 new
        else:
            metrics[key] = acc_class.get_reduced_value_from_states([state])

    for backend in self.logger_backends:
        if metrics:
            await backend.log(metrics, step)
        if samples:
            await backend.log_samples(samples, step)   # 👈 new hook

• With SAMPLE: flush will also collect dicts from SampleAccumulator.
• Backends get a new hook log_samples(samples, step) that handles rendering:
  • Console → print JSON.
  • WandB → create/update a wandb.Table.

---

Stage 2: User-Defined Filters

We can extend sample logging with user-defined filters to control what and how much gets logged. Logging every sample is often too expensive, and users may want to focus only on specific cases.

Examples of filters to expose:

• Frequency: log every Nth sample, or random subsampling.
• Reward thresholds: log only when reward ≤ 0 or ≥ X.
• Advantage thresholds: log only extreme positive/negative advantages.
• By step/policy version: log only at certain intervals.

User-facing config (conceptual):

sample_logging:
  enabled: true
  filters:
    reward_below: 0.1
    freq: 0.05     # log 5% of samples
    step: 100      # log every 100 steps

This would give users flexible, configurable control while keeping Stage 1 simple and minimal.

[DNXie]

cc @allenwang28 @felipemello1 @joecummings

[felipemello1]

looks good! :)

1. I chatted with @init27 and he mentioned that he wanted to see samples with the lowest reward. I am thinking that users will want different filtering logic, e.g by length.

We could do something like this instead:

metric_logging:
	sample_logging
		filter_fn: 
			_target_: path.to.filter_fn
			arg1:
			arg2:
	backends:
		wandb: ...
		console: ...

def my_filter_fn(groups: List[Group]) -> None:
	# do something
    record_metric(...)

If we go down that route, you could use _get_component_from_path to instantiate path.to.filter_fn.

I know that @allenwang28 wanted the config to be lightweight, so i am fine if we hardcode to log the lowest/highest reward on every batch. No configs. What do you guys think?

@init27 can you weight in from an UX perspective?
a) Would it be enough/too much to have highest/lowest total reward per batch per train step? Do we need a filter_fn?
b) Whats a good frequency to hardcode? Per train step, per group rollout, every N train steps? Should it be configurable?
c) Do we need to break it down per reward fn, or is it fine to have it per total_reward? e.g. code rewards always < creative writing, so i never log for creative writing

2. @DNXie , we have 2 backend.log places, so we need to make changes to both.
   a) in the MetricCollector, for local backends (i.e. we log per rank)
   b) In the global actor, for when we do reduce across ranks.

I would:

• See if it makes sense to use reduce_metrics_states for both and put the if reduction == Reduce.SAMPLE.value: logic there
• IMO, we can do the rest like you suggested. I would just consider the case where the user has a Reduce.SampleV2, or a reduce.HISTOGRAM. How to minimize touching N functions with multiple if/else. But just food for though. For now, its good.

[allenwang28]

I know that @allenwang28 wanted the config to be lightweight, so i am fine if we hardcode to log the lowest/highest reward on every batch. No configs. What do you guys think?

sgtm, I think lowest/highest + maybe a random sample. I would avoid having another config for this

[joecummings]

Couldn't we do all the "sampling" in post processing? Like log everything and then sort by lowest score. I fear that adding it to the config would be too messy and not needed for p0.

[felipemello1]

Like log everything and then sort by lowest score

do you mean log every single training sample to wandb? That sounds a bit much to me.

[joecummings]

Like log everything and then sort by lowest score

do you mean log every single training sample to wandb? That sounds a bit much to me.

Bit much how? It's certainly less to worry about in your config and your code.

Alternatively, I'd be fine defaulting to logging a random selection or high/low automatically and NOT exposing any filter function right now.

[felipemello1]

Bit much how?

e.g. logging 1M samples to wandb haha

Alternatively, I'd be fine defaulting to logging a random selection or high/low automatically and NOT exposing any filter function right now.

sg, it seems that we are going this route

[DNXie]

sgtm, I think lowest/highest + maybe a random sample. I would avoid having another config for this

sounds good! For now I will hardcode them in a filter function. And later we can consider supporting customized filter functions like @felipemello1 suggested maybe after PTC. I will implement it like this:

def filter_sample_fn(random_ratio=0.05, reward_lt: float=0.2, reward_gt: float=0.8)

to sample the 5% random episodes and the rest with <=0.2 reward, >=0.8 reward.

---

@joecummings I agree with Felipe that it is too much to log everything. Let's hardcode a filter function for now and see if we want to support customized functions in the future.

Is there any suggestions on the hyper parameters here?

[felipemello1]

Regarding this hardcoded filter, some things to think about

to sample the 5% random episodes and the rest with <=0.2 reward, >=0.8 reward.

1. Can we set a specific number, e.g. top 1, bottom 1, instead of making it variable?
2. Will the table contain reward value per reward fn or just total?
3. How many samples do we expect to see per train_step? As many rollouts as possible or fixed N bottom and N top?

[DNXie]

@felipemello1

Can we set a specific number, e.g. top 1, bottom 1, instead of making it variable?

Is this per rank or globally? It would be more straightforward and cheap to implement on rank-level. If on rank-level, we can maintain a heap in the accumulator and keep the top and bottom N samples. If we want to keep global top/bottom N, it would require us to store much more samples and filter at global flush.

Will the table contain reward value per reward fn or just total?

Let's track each reward and a total value

How many samples do we expect to see per train_step? As many rollouts as possible or fixed N bottom and N top?

We can support both filter by rate or by number. Here is my thoughts:

We can have a interface for filters:

class SampleFilter():
    def filter_append(self, sample: Dict) -> bool:
        """Return True if the sample should be kept at append time."""
        ...

    def filter_flush(self, samples: List[Dict]) -> List[Dict]:
        """Optional: further filter the collected samples at flush time."""
        return samples

Some examples of implementation:

Random ratio filter:

class RandomRatioFilter:
    def __init__(self, ratio=0.05):
        self.ratio = ratio

    def filter_append(self, sample):
        return random.random() < self.ratio

Filter by reward value:

class RewardThresholdFilter:
    def __init__(self, lt=None, gt=None):
        self.lt = lt
        self.gt = gt

    def filter_append(self, sample):
        r = sample.get("reward", 0.0)
        if self.lt is not None and r >= self.lt:
            return False
        if self.gt is not None and r <= self.gt:
            return False
        return True

This can also be customized to filter by certain reward if we pass all the reward value to record_metric

record_metric(
    "reward/samples",
    {"episode_id": ep.episode_id, "prompt": prompt, "reward": {"math": r1, "...": r2}},
    Reduce.SAMPLE,
)

Top bottom filter

import heapq

class TopBottomKFilter:
    def __init__(self, top_k=1, bottom_k=1, key="reward"):
        self.top_k = top_k
        self.bottom_k = bottom_k
        self.key = key
        self._top_heap = []     # min-heap for top-k
        self._bottom_heap = []  # max-heap for bottom-k (store -value)

    def filter_append(self, sample):
        # maintain two heaps

    def filter_flush(self, samples):
        # return the samples in both heaps

How does it sound?

[felipemello1]

Is this per rank or globally? It would be more straightforward and cheap to implement on rank-level.

Good question. The logic for all metrics is to accumulate per rank, and if you need to reduce across ranks, you use SampleAccumulator. get_reduced_value_from_states.

Let's track each reward and a total value

I prefer to log all rewards, but you may need to change some functions if you track each reward. I wonder if we should have this per-fn reward stored in the Group. But we can have that as a follow up. If things get too complex, lets just start with total reward.

We can have a interface for filters:

Given the feedback here to not support configs just yet, i think that having an interface / heap / filter variants at this moment is an overkill. How about we start with a simple function. Would that work?

def record_worst_and_best_sample(groups: List[Group]) -> None:
	best_trajectory = None
    worst_trajectory = None
    for group in gorups:
          for traj in group.trajectories:
                   if traj.total_reward > best_trajectory.reward:
                          best_trajectory = traj
				    if traj.total_reward < worst_trajectory.reward:
                          worst_trajectory = traj
     
      record_metrics("blabla/best_trajectory, best_trajectory, Reduce.Sample)
      record_metrics("blabla/worst_trajectory, worst_trajectory, Reduce.Sample)

Do we really need a class to hold a state? Its fine if the answer is yes. In that case, then the heap makes sense.

---

In Reduce.Sample.get_reduced_value_from_states, we just merge all (best, worst) * N_ranks. Wdyt?

[DNXie]

@felipemello1

Thanks for the suggestion! Regarding record_worst_and_best_sample

I see your point. The record_worst_and_best_sample function is simple and cheap. But I worry that this might already be too restrictive even at the experimental stage. In practice we often want more than just the absolute best and worst samples. So let's implement with a heap. And I’d argue that defining a filter interface is still worthwhile, as it gives us much more flexibility even for internal debugging purposes

---

Good question. The logic for all metrics is to accumulate per rank, and if you need to reduce across ranks, you use SampleAccumulator. get_reduced_value_from_states.

One open question is how we want to treat the “top/bottom” globally:

• Option A: each rank keeps its top-N and bottom-N locally, and get_reduced_value_from_states just merges those (simple, cheap).
• Option B: after merging, we apply the filter again globally to get the true global top/bottom-N (more precise, slightly more overhead and changes to the code).

Both are feasible. Personally, I’d prefer option A as it is more straightforward.