resolve comment

Author: DNXie

apps/grpo/main.py

@@ -29,7 +29,7 @@
 from forge.data.rewards import MathReward, ThinkingReward
 from forge.data_models.completion import Completion
 from forge.observability.metric_actors import get_or_create_metric_logger
-from forge.observability.metrics import record_episode_sample, record_metric, Reduce
+from forge.observability.metrics import record_metric, Reduce
 from forge.observability.perf_tracker import Tracer
 
 from forge.types import LauncherConfig, ProvisionerConfig
@@ -449,8 +449,13 @@ async def continuous_rollouts():
             for episode, advantage in zip(episodes, advantages):
                 episode.advantage = advantage
                 await replay_buffer.add.call_one(episode)
-                record_episode_sample(
-                    "main_samples/continuous_rollouts/sample_table", episode
+
+                sample = episode.to_dict(exclude=["ref_logprobs", "completion"])
+                sample["score"] = sample["reward"]
+                record_metric(
+                    "main_samples/continuous_rollouts/sample_table",
+                    sample,
+                    Reduce.SAMPLE,
                 )
 
             rollout_count += 1

src/forge/observability/__init__.py

@@ -21,7 +21,6 @@
     MetricAccumulator,
     MetricCollector,
     MinAccumulator,
-    record_episode_sample,
     record_metric,
     Reduce,
     reduce_metrics_states,
@@ -37,7 +36,6 @@
     # Main API functions
     "record_metric",
     "reduce_metrics_states",
-    "record_episode_sample",
     "get_logger_backend_class",
     "get_or_create_metric_logger",
     # Performance tracking

src/forge/observability/metrics.py

@@ -201,17 +201,6 @@ def reduce_metrics_states(states: list[dict[str, dict[str, Any]]]) -> list[Metri
     return reduced_metrics
 
 
-def record_episode_sample(table_name: str, episode):
-    """
-    Record a structured sample-level log for a single episode.
-    Args:
-        table_name (str): logging prefix (e.g. "rollout/sample").
-        episode (Episode): episode object with filled attributes.
-    """
-    sample = episode.to_dict(exclude=["ref_logprobs", "completion"])
-    record_metric(table_name, sample, Reduce.SAMPLE)
-
-
 ################
 # Accumulators #
 ################
@@ -430,7 +419,7 @@ class SampleAccumulator(MetricAccumulator):
     """
 
     def __init__(
-        self, reduction: Reduce, top_k: int = 1, bottom_k: int = 1, key: str = "reward"
+        self, reduction: Reduce, top_k: int = 1, bottom_k: int = 1, key: str = "score"
     ):
         super().__init__(reduction)
         self.samples: List[Dict[str, Any]] = []
@@ -869,12 +858,10 @@ async def finish(self) -> None:
     async def log_samples(self, samples: List[Metric], step: int) -> None:
         """Pretty-print sample-level logs to console."""
 
-        logger.info(f"==========  SAMPLE LOGS STEP {step} ==========")
         for sample in samples:
             table_name, table_rows = sample.key, sample.value
             logger.info(f"[{table_name}] ({len(table_rows)} samples)")
             logger.info(json.dumps(table_rows, indent=2, ensure_ascii=False))
-        logger.info("==============================================\n")
 
 
 class WandbBackend(LoggerBackend):
@@ -1056,6 +1043,13 @@ async def log_samples(self, samples: List[Metric], step: int) -> None:
 
             # Add rows (fill missing columns with None)
             for s in table_rows:
+                # Check for extra columns not in the table schema
+                extra_columns = set(s.keys()) - set(table.columns)
+                if extra_columns:
+                    logger.warning(
+                        f"WandbBackend: Row has extra columns not in table '{table_name}': {sorted(extra_columns)}. "
+                        f"These will be ignored."
+                    )
                 values = [s.get(c) for c in table.columns]
                 table.add_data(*values)