metrics(async): report selected policy provenance

docs/user-guide/fully-async.md

@@ -308,13 +308,33 @@ rollout/fully_async/queue_size
 rollout/fully_async/aborted_groups_filtered
 rollout/fully_async/stale_groups_filtered
 rollout/fully_async/avg_staleness, rollout/fully_async/max_staleness
+rollout/fully_async/avg_post_generation_staleness, rollout/fully_async/max_post_generation_staleness
+rollout/fully_async/avg_generation_version_span, rollout/fully_async/max_generation_version_span
+rollout/fully_async/token_weighted_staleness
+rollout/fully_async/weight_version_sample_coverage
 rollout/fully_async/buffer_avg_staleness, rollout/fully_async/buffer_max_staleness
 rollout/dynamic_filter/drop_<reason>
 ```
 
-The `avg_staleness` and `max_staleness` pair covers the groups training actually
-consumed, while the `buffer_` pair covers the groups still sitting in the buffer when
-the step drained it.
+Staleness is measured in published rollout weight versions: with
+`--update-weights-interval 1`, one version normally corresponds to one optimizer step;
+otherwise it does not. The metrics separate three parts of policy provenance:
+
+- `avg_staleness` and `max_staleness` are the gap from the current version to the
+  oldest version in each selected group. `stale_groups_filtered` counts groups that
+  exceeded `--max-weight-staleness`; rejected groups do not enter these averages.
+- `post_generation_staleness` is the gap from the current version to the newest version
+  in each selected group. `generation_version_span` is the newest minus oldest version
+  within the group. Together they distinguish aging after generation from updates that
+  landed while a trajectory was being generated.
+- `token_weighted_staleness` averages `current_version - span.version`, weighting each
+  numeric version span by its number of generated tokens. It therefore describes the
+  typical version gap of stamped generation tokens rather than the oldest token only.
+
+`weight_version_sample_coverage` is the fraction of selected samples with at least one
+numeric version span. Missing provenance is excluded from lag averages, not treated as
+zero. The `buffer_` pair covers groups still sitting in the buffer when the step drained
+it.
 
 A `queue_size` pinned at zero means rollout is the bottleneck, so scale rollout capacity
 or lower per-sample generation cost. A `queue_size` pinned at capacity means training is

miles/rollout/filter_hub/common_filters.py

@@ -1,4 +1,5 @@
 from argparse import Namespace
+from dataclasses import dataclass
 
 import torch
 
@@ -8,6 +9,64 @@
 Group = list[Sample | list[Sample]]
 
 
+@dataclass(frozen=True)
+class GroupWeightVersionStats:
+    sample_count: int
+    versioned_sample_count: int
+    versioned_token_count: int
+    token_version_sum: int
+    oldest_version: int | None
+    newest_version: int | None
+
+    def oldest_lag(self, current_version: int | None) -> int | None:
+        if self.oldest_version is None or current_version is None:
+            return None
+        return current_version - self.oldest_version
+
+    def newest_lag(self, current_version: int | None) -> int | None:
+        if self.newest_version is None or current_version is None:
+            return None
+        return current_version - self.newest_version
+
+    def token_weighted_lag(self, current_version: int | None) -> float | None:
+        if current_version is None or self.versioned_token_count == 0:
+            return None
+        return current_version - self.token_version_sum / self.versioned_token_count
+
+
+def group_weight_version_stats(group: Group) -> GroupWeightVersionStats:
+    sample_count = 0
+    versioned_sample_count = 0
+    versioned_token_count = 0
+    token_version_sum = 0
+    oldest_version = None
+    newest_version = None
+
+    for sample in iter_samples(group):
+        sample_count += 1
+        sample_has_version = False
+        for span in sample.all_weight_version_spans:
+            if not str(span.version).isdigit():
+                continue
+            version = int(span.version)
+            token_count = span.abs_end - span.abs_start
+            versioned_token_count += token_count
+            token_version_sum += version * token_count
+            oldest_version = version if oldest_version is None else min(oldest_version, version)
+            newest_version = version if newest_version is None else max(newest_version, version)
+            sample_has_version = True
+        versioned_sample_count += int(sample_has_version)
+
+    return GroupWeightVersionStats(
+        sample_count=sample_count,
+        versioned_sample_count=versioned_sample_count,
+        versioned_token_count=versioned_token_count,
+        token_version_sum=token_version_sum,
+        oldest_version=oldest_version,
+        newest_version=newest_version,
+    )
+
+
 def apply_preput_filters(args: Namespace, dynamic_filter, samples: Group, **kwargs) -> FilterOutput:
     output = apply_aborted_filter(args, samples, **kwargs)
     if not output.keep:
@@ -43,8 +102,4 @@ def apply_reward_nonzero_std_filter(args, samples: list[Sample | list[Sample]],
 
 
 def group_staleness(group: Group, current_version: int | None) -> int | None:
-    versions = [version for sample in iter_samples(group) if (version := sample.oldest_weight_version) is not None]
-    oldest = min(versions) if versions else None
-    if oldest is None or current_version is None:
-        return None
-    return current_version - oldest
+    return group_weight_version_stats(group).oldest_lag(current_version)

miles/rollout/fully_async_data_buffer.py

@@ -15,7 +15,13 @@
 from dataclasses import dataclass
 
 from miles.rollout.filter_hub.base_types import MetricGatherer, call_dynamic_filter
-from miles.rollout.filter_hub.common_filters import apply_aborted_filter, apply_missing_reward_filter, group_staleness
+from miles.rollout.filter_hub.common_filters import (
+    GroupWeightVersionStats,
+    apply_aborted_filter,
+    apply_missing_reward_filter,
+    group_staleness,
+    group_weight_version_stats,
+)
 from miles.utils.function_registry import load_function
 from miles.utils.types import Sample
 
@@ -109,6 +115,12 @@ def __init__(self, input: DataBufferConstructorInput):
         self._metric_aborted_groups = 0
         self._metric_stale_groups = 0
         self._metric_consumed_staleness: list[int] = []
+        self._metric_selected_newest_lag: list[int] = []
+        self._metric_selected_version_span: list[int] = []
+        self._metric_selected_token_lag_sum = 0.0
+        self._metric_selected_versioned_tokens = 0
+        self._metric_selected_samples = 0
+        self._metric_selected_versioned_samples = 0
 
     async def put(self, input: DataBufferInput) -> None:
         if not self._preput_filter(input):
@@ -148,16 +160,38 @@ async def get(self, current_version: int | None = None, **_) -> DataBufferInput:
                 entry = self._buffer.pop(0)
                 self._cond.notify_all()  # wake producers blocked on a full buffer
 
-                staleness = group_staleness(entry.group, current_version)
+                version_stats = group_weight_version_stats(entry.group)
+                staleness = version_stats.oldest_lag(current_version)
                 if staleness is not None:
-                    self._metric_consumed_staleness.append(staleness)
                     if self._args.max_weight_staleness is not None and staleness > self._args.max_weight_staleness:
                         logger.info(f"Filtered stale group ({staleness=} > max={self._args.max_weight_staleness})")
                         self._metric_stale_groups += 1
                         self._unused_handler_fn(entry.prompt_group)
                         continue
+                    self._metric_consumed_staleness.append(staleness)
+                self._record_selected_version_stats(version_stats, current_version)
                 return entry
 
+    def _record_selected_version_stats(
+        self,
+        stats: GroupWeightVersionStats,
+        current_version: int | None,
+    ) -> None:
+        self._metric_selected_samples += stats.sample_count
+        self._metric_selected_versioned_samples += stats.versioned_sample_count
+
+        if stats.oldest_version is not None and stats.newest_version is not None:
+            self._metric_selected_version_span.append(stats.newest_version - stats.oldest_version)
+
+        newest_lag = stats.newest_lag(current_version)
+        if newest_lag is not None:
+            self._metric_selected_newest_lag.append(newest_lag)
+
+        token_weighted_lag = stats.token_weighted_lag(current_version)
+        if token_weighted_lag is not None:
+            self._metric_selected_token_lag_sum += token_weighted_lag * stats.versioned_token_count
+            self._metric_selected_versioned_tokens += stats.versioned_token_count
+
     def get_metrics(self) -> dict[str, float]:
         prefix = "rollout/fully_async/"
         metrics = {
@@ -169,6 +203,20 @@ def get_metrics(self) -> dict[str, float]:
         if consumed := self._metric_consumed_staleness:
             metrics[f"{prefix}avg_staleness"] = sum(consumed) / len(consumed)
             metrics[f"{prefix}max_staleness"] = max(consumed)
+        if newest_lag := self._metric_selected_newest_lag:
+            metrics[f"{prefix}avg_post_generation_staleness"] = sum(newest_lag) / len(newest_lag)
+            metrics[f"{prefix}max_post_generation_staleness"] = max(newest_lag)
+        if version_span := self._metric_selected_version_span:
+            metrics[f"{prefix}avg_generation_version_span"] = sum(version_span) / len(version_span)
+            metrics[f"{prefix}max_generation_version_span"] = max(version_span)
+        if self._metric_selected_versioned_tokens:
+            metrics[f"{prefix}token_weighted_staleness"] = (
+                self._metric_selected_token_lag_sum / self._metric_selected_versioned_tokens
+            )
+        if self._metric_selected_samples:
+            metrics[f"{prefix}weight_version_sample_coverage"] = (
+                self._metric_selected_versioned_samples / self._metric_selected_samples
+            )
         buffered = [
             s for entry in self._buffer if (s := group_staleness(entry.group, self._current_version)) is not None
         ]
@@ -178,5 +226,11 @@ def get_metrics(self) -> dict[str, float]:
 
         self._metric_gatherer = MetricGatherer()
         self._metric_consumed_staleness = []
+        self._metric_selected_newest_lag = []
+        self._metric_selected_version_span = []
+        self._metric_selected_token_lag_sum = 0.0
+        self._metric_selected_versioned_tokens = 0
+        self._metric_selected_samples = 0
+        self._metric_selected_versioned_samples = 0
         self._metric_aborted_groups = self._metric_stale_groups = 0
         return metrics

tests/fast/rollout/test_filters.py

@@ -15,6 +15,7 @@
     apply_preput_filters,
     apply_reward_nonzero_std_filter,
     group_staleness,
+    group_weight_version_stats,
 )
 from miles.utils.function_registry import load_function
 from miles.utils.types import Sample, WeightVersionSpan, WeightVersionsPerCall
@@ -248,3 +249,27 @@ def test_group_staleness_uses_oldest_version_across_nested_samples():
     assert group_staleness(group, current_version=None) is None
     assert group_staleness([make_sample()], current_version=10) is None
     assert group_staleness([make_sample(weight_versions=("12",))], current_version=10) == -2
+
+
+def test_group_weight_version_stats_preserve_sample_and_token_populations():
+    first = make_sample()
+    first.weight_versions = [
+        WeightVersionsPerCall(
+            spans=[
+                WeightVersionSpan(version="4", abs_start=2, abs_end=5),
+                WeightVersionSpan(version="6", abs_start=5, abs_end=6),
+            ]
+        )
+    ]
+    second = make_sample(weight_versions=("not-numeric",))
+
+    stats = group_weight_version_stats([first, [second]])
+
+    assert stats.sample_count == 2
+    assert stats.versioned_sample_count == 1
+    assert stats.versioned_token_count == 4
+    assert stats.oldest_lag(current_version=10) == 6
+    assert stats.newest_lag(current_version=10) == 4
+    assert stats.token_weighted_lag(current_version=10) == 5.5
+    assert stats.oldest_lag(current_version=None) is None
+    assert stats.token_weighted_lag(current_version=None) is None

tests/fast/rollout/test_fully_async_rollout.py

@@ -256,7 +256,7 @@ def get_samples_with_fresh_versions(num_samples):
 
     assert data_source.recycled == [stale]
     assert output.metrics["rollout/fully_async/stale_groups_filtered"] == 1
-    assert output.metrics["rollout/fully_async/max_staleness"] == 5
+    assert output.metrics["rollout/fully_async/max_staleness"] == 0
 
 
 async def test_stale_group_dropped_by_default(monkeypatch):
@@ -587,6 +587,51 @@ async def test_buffer_staleness_metrics():
     assert metrics["rollout/fully_async/buffer_max_staleness"] == 4
 
 
+async def test_buffer_reports_selected_policy_provenance():
+    buffer, _ = make_buffer(max_groups=8, max_staleness=2)
+    await put_group(buffer, make_group(1, weight_versions=["2", "4"]))
+    await put_group(buffer, make_group(2, weight_versions=["8", "10"]))
+
+    assert (await buffer.get(current_version=10)).group[0].group_index == 2
+    metrics = buffer.get_metrics()
+
+    assert metrics["rollout/fully_async/avg_staleness"] == 2
+    assert metrics["rollout/fully_async/max_staleness"] == 2
+    assert metrics["rollout/fully_async/avg_post_generation_staleness"] == 0
+    assert metrics["rollout/fully_async/avg_generation_version_span"] == 2
+    assert metrics["rollout/fully_async/token_weighted_staleness"] == 1
+    assert metrics["rollout/fully_async/weight_version_sample_coverage"] == 1
+
+
+async def test_buffer_reports_generation_span_without_current_version():
+    buffer, _ = make_buffer(max_groups=8)
+    await put_group(buffer, make_group(1, weight_versions=["2", "4"]))
+
+    await buffer.get(current_version=None)
+    metrics = buffer.get_metrics()
+
+    assert metrics["rollout/fully_async/avg_generation_version_span"] == 2
+    assert "rollout/fully_async/avg_post_generation_staleness" not in metrics
+    assert "rollout/fully_async/token_weighted_staleness" not in metrics
+
+
+async def test_buffer_reports_missing_weight_version_coverage_without_inventing_lag():
+    buffer, _ = make_buffer(max_groups=8)
+    partial_group = make_group(1, weight_versions=["7"])
+    partial_group[1].weight_versions = []
+    await put_group(buffer, partial_group)
+
+    await buffer.get(current_version=10)
+    metrics = buffer.get_metrics()
+
+    assert metrics["rollout/fully_async/weight_version_sample_coverage"] == 0.5
+    assert metrics["rollout/fully_async/token_weighted_staleness"] == 3
+
+    metrics = buffer.get_metrics()
+    assert "rollout/fully_async/weight_version_sample_coverage" not in metrics
+    assert "rollout/fully_async/token_weighted_staleness" not in metrics
+
+
 class RecordingBuffer(data_buffer.DefaultDataBuffer):
     constructed_with = None