[RLLib] Episode media logging support

rllib/evaluation/episode.py

@@ -67,6 +67,7 @@ def __init__(self, policies: Dict[PolicyID, Policy],
         self.custom_metrics: Dict[str, float] = {}
         self.user_data: Dict[str, Any] = {}
         self.hist_data: Dict[str, List[float]] = {}
+        self.media: Dict[str, Any] = {}
         self._policies: Dict[PolicyID, Policy] = policies
         self._policy_mapping_fn: Callable[[AgentID], PolicyID] = \
             policy_mapping_fn

rllib/evaluation/metrics.py

@@ -122,6 +122,8 @@ def summarize_episodes(
     custom_metrics = collections.defaultdict(list)
     perf_stats = collections.defaultdict(list)
     hist_stats = collections.defaultdict(list)
+    episode_media = collections.defaultdict(list)
+
     for episode in episodes:
         episode_lengths.append(episode.episode_length)
         episode_rewards.append(episode.episode_reward)
@@ -134,6 +136,8 @@ def summarize_episodes(
                 policy_rewards[policy_id].append(reward)
         for k, v in episode.hist_data.items():
             hist_stats[k] += v
+        for k, v in episode.media.items():
+            episode_media[k].append(v)
     if episode_rewards:
         min_reward = min(episode_rewards)
         max_reward = max(episode_rewards)
@@ -191,6 +195,7 @@ def summarize_episodes(
         episode_reward_min=min_reward,
         episode_reward_mean=avg_reward,
         episode_len_mean=avg_length,
+        episode_media=dict(episode_media),
         episodes_this_iter=len(new_episodes),
         policy_reward_min=policy_reward_min,
         policy_reward_max=policy_reward_max,

rllib/evaluation/rollout_metrics.py

@@ -2,7 +2,12 @@
 
 # Define this in its own file, see #5125
 RolloutMetrics = collections.namedtuple("RolloutMetrics", [
-    "episode_length", "episode_reward", "agent_rewards", "custom_metrics",
-    "perf_stats", "hist_data"
+    "episode_length",
+    "episode_reward",
+    "agent_rewards",
+    "custom_metrics",
+    "perf_stats",
+    "hist_data",
+    "media",
 ])
-RolloutMetrics.__new__.__defaults__ = (0, 0, {}, {}, {}, {})
+RolloutMetrics.__new__.__defaults__ = (0, 0, {}, {}, {}, {}, {})

rllib/evaluation/rollout_worker.py

@@ -921,6 +921,7 @@ def get_metrics(self) -> List[Union[RolloutMetrics, OffPolicyEstimate]]:
         # Get metrics from our reward-estimators (if any).
         for m in self.reward_estimators:
             out.extend(m.get_metrics())
+
         return out
 
     @DeveloperAPI

rllib/evaluation/sampler.py

@@ -827,7 +827,7 @@ def _process_observations(
                     RolloutMetrics(episode.length, episode.total_reward,
                                    dict(episode.agent_rewards),
                                    episode.custom_metrics, {},
-                                   episode.hist_data))
+                                   episode.hist_data, episode.media))
         else:
             hit_horizon = False
             all_agents_done = False
@@ -1050,7 +1050,7 @@ def _process_observations_w_trajectory_view_api(
                     RolloutMetrics(episode.length, episode.total_reward,
                                    dict(episode.agent_rewards),
                                    episode.custom_metrics, {},
-                                   episode.hist_data))
+                                   episode.hist_data, episode.media))
         else:
             hit_horizon = False
             all_agents_done = False

rllib/execution/metric_ops.py

@@ -108,6 +108,7 @@ def __call__(self, _: Any) -> Dict:
         res["info"] = info
         res["info"].update(counters)
         res["custom_metrics"] = res.get("custom_metrics", {})
+        res["episode_media"] = res.get("episode_media", {})
         res["custom_metrics"].update(custom_metrics_from_info)
         return res