From 922a1af76e860eb34ce5e0b979fad4dd3f6a74e7 Mon Sep 17 00:00:00 2001
From: Jun Gong <jungong@anyscale.com>
Date: Thu, 10 Nov 2022 12:55:34 -0800
Subject: [PATCH] [RLlib] minor bug fix. Resetted exceptions are saved under
 env_id key. (#30183)

Signed-off-by: Jun Gong <jungong@anyscale.com>
Signed-off-by: Weichen Xu <weichen.xu@databricks.com>
---
 rllib/evaluation/env_runner_v2.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/rllib/evaluation/env_runner_v2.py b/rllib/evaluation/env_runner_v2.py
index c14535b41de64..2893f8735dde1 100644
--- a/rllib/evaluation/env_runner_v2.py
+++ b/rllib/evaluation/env_runner_v2.py
@@ -798,7 +798,9 @@ def _handle_done_episode(
                 resetted_obs: Dict[
                     EnvID, Dict[AgentID, EnvObsType]
                 ] = self._base_env.try_reset(env_id)
-                if resetted_obs is None or not isinstance(resetted_obs, Exception):
+                if resetted_obs is None or not isinstance(
+                    resetted_obs[env_id], Exception
+                ):
                     break
                 else:
                     # Report a faulty episode.