move avg reward out

Author: DNXie

apps/grpo/main.py

@@ -148,7 +148,7 @@ class RewardActor(ForgeActor):
     @endpoint
     async def evaluate_response(
         self, prompt: str, response: str, target: str
-    ) -> dict[str, float]:
+    ) -> (dict[str, float], float):
         total_rewards = 0.0
         reward_breakdown = {}  # reward breakdown by function
         for reward_fn in self.reward_functions:
@@ -189,9 +189,8 @@ async def evaluate_response(
                 Reduce.SUM,
             )
 
-        avg_reward = total_rewards / len(self.reward_functions)
-        reward_breakdown["reward"] = avg_reward
-        return reward_breakdown
+        avg_reward: float = total_rewards / len(self.reward_functions)
+        return reward_breakdown, avg_reward
 
 
 @dataclass
@@ -397,10 +396,12 @@ async def continuous_rollouts():
                     response=response.text,
                     completion=response,
                 )
-                episode.reward_breakdown = await reward_actor.evaluate_response.route(
+                (
+                    episode.reward_breakdown,
+                    episode.reward,
+                ) = await reward_actor.evaluate_response.route(
                     prompt=prompt, response=response.text, target=target
                 )
-                episode.reward = episode.reward_breakdown["reward"]
                 episodes.append(episode)
 
                 # Build input_ids for reference logprobs

src/forge/observability/metrics.py

@@ -223,6 +223,7 @@ def record_episode_sample(table_name: str, episode):
         **(
             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,