add to_dict

Author: DNXie

apps/grpo/main.py

@@ -78,6 +78,30 @@ def response_tensor(self) -> torch.Tensor:
             tensor = F.pad(tensor, (0, diff), value=self.pad_id)
         return tensor
 
+    def to_dict(self, exclude: list[str] | None = None) -> dict[str, Any]:
+        """Convert episode to dict, optionally excluding specified fields."""
+        result = {
+            "episode_id": self.episode_id,
+            "policy_version": self.policy_version,
+            "prompt": self.request,
+            "response": self.response,
+            "target": str(self.target),
+            "reward": self.reward,
+            "advantage": self.advantage,
+            "request_len": self.request_len,
+            "response_len": self.response_len,
+            "pad_id": self.pad_id,
+        }
+
+        if self.reward_breakdown is not None:
+            result.update(self.reward_breakdown)
+
+        if exclude:
+            for key in exclude:
+                result.pop(key, None)
+
+        return result
+
 
 # Represents the group (G) of episodes in GRPO
 Group = list[Episode]

src/forge/observability/metrics.py

@@ -194,21 +194,7 @@ def record_episode_sample(table_name: str, episode):
         table_name (str): logging prefix (e.g. "rollout/sample").
         episode (Episode): episode object with filled attributes.
     """
-    sample = {
-        "episode_id": episode.episode_id,
-        "policy_version": episode.policy_version,
-        "prompt": episode.request,
-        "response": episode.response,
-        "target": str(episode.target),
-        **(
-            episode.reward_breakdown or {}
-        ),  # per-fn breakdown including the average reward
-        "reward": episode.reward,
-        "advantage": episode.advantage,
-        "request_len": episode.request_len,
-        "response_len": episode.response_len,
-        "pad_id": episode.pad_id,
-    }
+    sample = episode.to_dict()
     record_metric(table_name, sample, Reduce.SAMPLE)