Log compaction-aware rollout metrics

miles/ray/rollout/metrics.py

@@ -96,6 +96,7 @@ def _compute_metrics_from_samples(args, samples):
     response_lengths = [sample.effective_response_length for sample in samples]
 
     log_dict = {}
+    log_dict |= _compute_training_sample_metrics(args, samples)
     log_dict |= dict_add_prefix(compute_statistics(response_lengths), "response_len/")
     log_dict |= _compute_zero_std_metrics(args, samples)
     log_dict |= _compute_spec_metrics(args, samples)
@@ -135,6 +136,34 @@ def _compute_metrics_from_samples(args, samples):
     return log_dict
 
 
+def _compute_training_sample_metrics(args: Any, samples: list[Sample]) -> dict[str, float | int]:
+    """Count training rows and average raw reward with equal weight per rollout.
+
+    Session compaction can turn one rollout into several training samples. The
+    sample count includes every resulting row, while the reward first averages
+    sibling rows that share a rollout ID so long rollouts do not receive more
+    metric weight merely because they produced more samples.
+    """
+    rewards_by_rollout: dict[tuple[str, int | None, int], list[float]] = {}
+    use_metadata_reward = bool(samples and samples[0].metadata and "raw_reward" in samples[0].metadata)
+    for position, sample in enumerate(samples):
+        if sample.rollout_id is not None:
+            rollout_key = ("rollout", sample.group_index, sample.rollout_id)
+        elif sample.index is not None:
+            rollout_key = ("sample", sample.group_index, sample.index)
+        else:
+            rollout_key = ("position", sample.group_index, position)
+
+        raw_reward = sample.metadata["raw_reward"] if use_metadata_reward else sample.get_reward_value(args)
+        rewards_by_rollout.setdefault(rollout_key, []).append(raw_reward)
+
+    rollout_rewards = [sum(rewards) / len(rewards) for rewards in rewards_by_rollout.values()]
+    return {
+        "num_training_samples": len(samples),
+        "episode_raw_reward": sum(rollout_rewards) / len(rollout_rewards) if rollout_rewards else 0.0,
+    }
+
+
 def _compute_perf_metrics_from_samples(args, samples, rollout_time):
     non_generation_time = [sample.non_generation_time for sample in samples]
 

tests/fast/ray/rollout/test_metrics.py

@@ -1,16 +1,76 @@
 from __future__ import annotations
 
 import pytest
-from tests.fast.ray.rollout.conftest import make_args, make_samples_grouped
+from tests.fast.ray.rollout.conftest import make_args, make_sample, make_samples_grouped
 
 from miles.ray.rollout.metrics import (
     _compute_metrics_from_samples,
     _compute_passrate_from_samples,
+    _compute_training_sample_metrics,
     _compute_zero_std_metrics,
     log_rollout_data,
 )
 
 
+class TestTrainingSampleMetrics:
+    def test_compacted_rollouts_are_counted_as_samples_but_rewarded_as_episodes(self):
+        args = make_args(reward_key=None)
+        samples = [
+            make_sample(index=0, rollout_id=10, reward=1.0),
+            make_sample(index=0, rollout_id=10, reward=1.0),
+            make_sample(index=0, rollout_id=10, reward=1.0),
+            make_sample(index=1, rollout_id=11, reward=0.0),
+        ]
+
+        out = _compute_training_sample_metrics(args, samples)
+
+        assert out["num_training_samples"] == 4
+        assert out["episode_raw_reward"] == pytest.approx(0.5)
+
+    def test_different_sibling_rewards_are_averaged_within_rollout_first(self):
+        args = make_args(reward_key=None)
+        samples = [
+            make_sample(index=0, rollout_id=10, reward=0.0),
+            make_sample(index=0, rollout_id=10, reward=1.0),
+            make_sample(index=1, rollout_id=11, reward=1.0),
+        ]
+
+        out = _compute_training_sample_metrics(args, samples)
+
+        assert out["episode_raw_reward"] == pytest.approx(0.75)
+
+    def test_rollout_ids_are_scoped_by_prompt_group(self):
+        args = make_args(reward_key=None)
+        samples = [
+            make_sample(group_index=0, index=0, rollout_id=10, reward=1.0),
+            make_sample(group_index=0, index=0, rollout_id=10, reward=1.0),
+            make_sample(group_index=1, index=1, rollout_id=10, reward=0.0),
+        ]
+
+        out = _compute_training_sample_metrics(args, samples)
+
+        assert out["episode_raw_reward"] == pytest.approx(0.5)
+
+    def test_metadata_raw_reward_and_fallback_identities(self):
+        args = make_args(reward_key=None)
+        samples = [
+            make_sample(index=5, reward=0.0),
+            make_sample(index=None, reward=0.0),
+        ]
+        samples[0].metadata = {"raw_reward": 1.0}
+        samples[1].metadata = {"raw_reward": 0.0}
+
+        out = _compute_training_sample_metrics(args, samples)
+
+        assert out == {"num_training_samples": 2, "episode_raw_reward": pytest.approx(0.5)}
+
+    def test_empty_samples(self):
+        assert _compute_training_sample_metrics(make_args(), []) == {
+            "num_training_samples": 0,
+            "episode_raw_reward": 0.0,
+        }
+
+
 class TestComputeZeroStdMetrics:
     def test_returns_empty_for_ppo_regardless_of_reward_distribution(self):
         args = make_args(advantage_estimator="ppo")
@@ -152,6 +212,8 @@ def test_rollout_log_fans_out_versioned_tito_keys(self, monkeypatch):
 
         log_rollout_data(0, args, samples, None, 1.0)
 
+        assert logged["rollout/num_training_samples"] == 4
+        assert logged["rollout/episode_raw_reward"] == pytest.approx(1.5)
         assert logged["rollout/tito_session_mismatch_rate/v2/assistant_text"] == 0.25
         assert "rollout/tito_session_mismatch_rate/assistant_text" not in logged