display reward in swarm overwatch

Author: binary-husky

ajet/task_rollout/native_parallel_worker.py

@@ -318,13 +318,17 @@ def stop_condition_callback(completed_task_id_map_ct):
         def update_rollout_result_array_preview(observation_window, completed_task_id_map_ct: Dict[str, List[SingleAgentContextTracker]]):
             buffer = ""
             completed_tasks_details = {}
+            completed_tasks_rewards = {}
             for task_id, tracker_arr in completed_task_id_map_ct.items():
                 buffer += f"Task {task_id} (completed {len(tracker_arr)} episodes):\n"
                 episode_uuids = []
+                rewards = []
                 for ct in tracker_arr:
                     buffer += f"\tEpisode: {ct.episode_uuid}\tTimelines: {len(ct.saved_timelines)}\tLLM_Calls: {ct.llm_call_cnt}\tReward: {ct.reward_structure.performance_reward}\n"
                     episode_uuids.append(ct.episode_uuid)
+                    rewards.append(float(ct.reward_structure.performance_reward))
                 completed_tasks_details[task_id] = episode_uuids
+                completed_tasks_rewards[task_id] = rewards
             buffer += f"\n"
             buffer += f"\n"
             counts = count_tasks(completed_task_id_map_ct)
@@ -345,6 +349,7 @@ def update_rollout_result_array_preview(observation_window, completed_task_id_ma
                 completed_non_dummy_task_target=n_batch_task,
                 task_expected_num_repeat=rollout_n,
                 completed_tasks_details=completed_tasks_details,
+                completed_tasks_rewards=completed_tasks_rewards,
             )
             http_update_rollout_pool_information(self.config, pool_info)
             return

ajet/tuner_lib/experimental/swarm_overwatch_utils.py

@@ -12,6 +12,7 @@ class CurrentBatchRolloutPoolInformation(BaseModel):
     completed_non_dummy_task_target: int = 0
     task_expected_num_repeat: int = 0
     completed_tasks_details: Dict[str, List[str]] = {}  # task_id -> list of episode_uuids
+    completed_tasks_rewards: Dict[str, List[float]] = {}  # task_id -> list of rewards (one per episode)
     running_episode_details: Dict[str, Dict[str, str]] | None = None # episode_uuid -> { "episode_status": ..., "time_since_last_activity": ..., "discard_episode_timeout": ..., "llm_call_count": ...}
     engine_status: str | None = None
     global_step: int | None = None

ajet/utils/swarm_overwatch.py

@@ -8,6 +8,7 @@
 from typing import Optional
 
 import httpx
+import numpy as np
 from rich.console import Console
 from rich.live import Live
 from rich.table import Table
@@ -296,10 +297,11 @@ def create_task_details_table(
 
         table.add_column("Task ID", style="cyan", no_wrap=True, overflow="ellipsis")
         table.add_column("Episodes", justify="right", style="green")
+        table.add_column("Reward", justify="right", style="yellow")
         table.add_column("Episode UUIDs (first 3)", style="dim", overflow="fold")
 
         if not info.completed_tasks_details:
-            table.add_row("[dim]No task details available[/dim]", "", "")
+            table.add_row("[dim]No task details available[/dim]", "", "", "")
             return table
 
         # Sort tasks by number of completed episodes (descending)
@@ -315,15 +317,25 @@ def create_task_details_table(
             if len(episode_uuids) > 3:
                 uuid_str += f" (+{len(episode_uuids) - 3} more)"
 
+            # Calculate reward statistics
+            reward_str = "-"
+            if info.completed_tasks_rewards and task_id in info.completed_tasks_rewards:
+                rewards = info.completed_tasks_rewards[task_id]
+                if rewards:
+                    mean_reward = np.mean(rewards)
+                    std_reward = np.std(rewards)
+                    reward_str = f"{mean_reward:.3f} ± {std_reward:.3f}"
+
             table.add_row(
                 task_id[:40] if len(task_id) > 40 else task_id,
                 f"{len(episode_uuids):,}",
+                reward_str,
                 uuid_str,
             )
 
         if len(sorted_tasks) > 30:
             table.add_row(
-                f"[dim]... and {len(sorted_tasks) - 30} more tasks[/dim]", "", ""
+                f"[dim]... and {len(sorted_tasks) - 30} more tasks[/dim]", "", "", ""
             )
 
         return table