integrate sampling

Author: DNXie

apps/grpo/main.py

@@ -30,7 +30,7 @@
 from forge.controller.provisioner import shutdown
 from forge.data.rewards import MathReward, ThinkingReward
 from forge.observability.metric_actors import get_or_create_metric_logger
-from forge.observability.metrics import record_metric, Reduce
+from forge.observability.metrics import record_episode_sample, record_metric, Reduce
 from forge.observability.perf_tracker import Tracer
 from forge.util.ops import compute_logprobs
 from monarch.actor import endpoint
@@ -54,6 +54,7 @@ class Episode:
     response_tokens: list[int] | None = None
     ref_logprobs: torch.Tensor | None = None
     reward: float | None = None
+    reward_breakdown: dict[str, float] | None = None
     advantage: float | None = None
 
     @property
@@ -168,8 +169,11 @@ class RewardActor(ForgeActor):
     reward_functions: list[Callable]
 
     @endpoint
-    async def evaluate_response(self, prompt: str, response: str, target: str) -> float:
+    async def evaluate_response(
+        self, prompt: str, response: str, target: str
+    ) -> dict[str, float]:
         total_rewards = 0.0
+        reward_breakdown = {}  # reward breakdown by function
         for reward_fn in self.reward_functions:
             reward = reward_fn(prompt, response, target)
             total_rewards += reward
@@ -178,6 +182,7 @@ async def evaluate_response(self, prompt: str, response: str, target: str) -> fl
             reward_fn_name = getattr(
                 reward_fn, "__name__", reward_fn.__class__.__name__
             )
+            reward_breakdown[reward_fn_name] = reward
             # per function reward
             record_metric(
                 f"reward/evaluate_response/sum_{reward_fn_name}_reward",
@@ -210,7 +215,8 @@ async def evaluate_response(self, prompt: str, response: str, target: str) -> fl
             )
 
         avg_reward = total_rewards / len(self.reward_functions)
-        return avg_reward
+        reward_breakdown["reward"] = avg_reward
+        return reward_breakdown
 
 
 @dataclass
@@ -395,9 +401,10 @@ async def continuous_rollouts():
                 episode.response = response.text
                 input_ids[i, :max_req_tokens] = episode.request_tensor
                 input_ids[i, max_req_tokens:] = episode.response_tensor
-                episode.reward = await reward_actor.evaluate_response.route(
+                episode.reward_breakdown = await reward_actor.evaluate_response.route(
                     prompt=prompt, response=response.text, target=target
                 )
+                episode.reward = episode.reward_breakdown["reward"]
 
             t.step("reward_evaluation")
 
@@ -416,7 +423,9 @@ async def continuous_rollouts():
             for episode, advantage in zip(group.episodes, advantages):
                 episode.advantage = advantage
                 await replay_buffer.add.call_one(episode)
+                record_episode_sample("rollout/sample", episode)
 
+            record_metric("sample/", {}, Reduce.SAMPLE)
             # Log metrics
             rollout_count += 1
             record_metric(

src/forge/observability/metrics.py

@@ -113,6 +113,38 @@ def record_metric(key: str, value: Any, reduction: Reduce = Reduce.MEAN) -> None
     collector.push(key, value, reduction)
 
 
+def record_episode_sample(key: str, episode):
+    """
+    Record a structured sample-level log for a single episode.
+
+    Args:
+        key (str): logging prefix (e.g. "rollout/sample").
+        episode (Episode): episode object with filled attributes.
+        reward_breakdown (dict[str, float]): per-function rewards, e.g. {"MathReward": 0.8, "FormatReward": 1.0}.
+    """
+    sample = {
+        "episode_id": episode.episode_id,
+        "policy_version": episode.policy_version,
+        "prompt": episode.request,
+        "response": episode.response,
+        "target": episode.target,
+        **(
+            episode.reward_breakdown or {}
+        ),  # per-fn breakdown including the average reward
+        "advantage": episode.advantage,
+        "ref_logprobs": (
+            episode.ref_logprobs.mean().item()
+            if episode.ref_logprobs is not None
+            else None
+        ),
+        "request_len": episode.request_len,
+        "response_len": episode.response_len,
+        "pad_id": episode.pad_id,
+    }
+
+    record_metric(key, sample, Reduce.SAMPLE)
+
+
 def reduce_metrics_states(states: List[Dict[str, Dict[str, Any]]]) -> Dict[str, Any]:
     """Reduce metric accumulators states to a single value per metric.
 
@@ -465,7 +497,9 @@ class SampleAccumulator(MetricAccumulator):
     Optionally uses a SampleFilter to decide what to keep at append/flush time.
     """
 
-    def __init__(self, reduction: Reduce, filter: SampleFilter | None = None):
+    def __init__(
+        self, reduction: Reduce, filter: SampleFilter | None = TopBottomKFilter()
+    ):
         super().__init__(reduction)
         self.samples: List[Dict[str, Any]] = []
         self.filter = filter
@@ -598,6 +632,7 @@ def push(self, key: str, value: Any, reduction: Reduce = Reduce.MEAN) -> None:
             raise ValueError("Collector not initialized—call init first")
 
         if key not in self.accumulators:
+            # TODO: make sample filter configurable
             self.accumulators[key] = reduction.accumulator_class(reduction)
 
         self.accumulators[key].append(value)
@@ -724,6 +759,20 @@ async def log(self, metrics: Dict[str, Any], step: int) -> None:
             logger.info(f"  {key}: {value}")
         logger.info("==============================\n")
 
+    async def log_samples(self, samples: Dict[str, List[dict]], step: int) -> None:
+        """Pretty-print sample-level logs to console."""
+        if not samples:
+            return
+        import pprint
+
+        logger.info(f"=== [{self.prefix}] - SAMPLE LOGS STEP {step} ===")
+        for key, rows in samples.items():
+            logger.info(f"[{key}] ({len(rows)} samples)")
+            for sample in rows:
+                pretty = pprint.pformat(sample, indent=4, width=120, compact=True)
+                logger.info(pretty)
+        logger.info("==============================================\n")
+
     async def finish(self) -> None:
         pass
 
@@ -836,6 +885,28 @@ async def log(self, metrics: Dict[str, Any], step: int) -> None:
         else:
             logger.debug(f"WandbBackend: No run started, skipping log for {self.name}")
 
+    async def log_samples(self, samples: Dict[str, List[dict]], step: int) -> None:
+        """Log sample-level data to WandB Tables."""
+        import wandb
+
+        if not self.run or not samples:
+            return
+
+        for key, rows in samples.items():
+            if not rows:
+                continue
+
+            # Create a WandB Table dynamically based on keys of first sample
+            columns = list(rows[0].keys())
+            table = wandb.Table(columns=columns)
+            for sample in rows:
+                table.add_data(*[sample.get(c) for c in columns])
+
+            self.run.log({f"{key}_table": table, "global_step": step})
+            logger.info(
+                f"WandbBackend: Logged {len(rows)} samples for {key} at step {step}"
+            )
+
     def get_metadata_for_secondary_ranks(self) -> Dict[str, Any]:
         if self.run and not self.reduce_across_ranks and self.share_run_id:
             return {"shared_run_id": self.run.id}