Record episode status.

examples/hitl/rearrange_v2/rearrange_v2.py

@@ -42,6 +42,7 @@
     GuiRobotController,
 )
 from habitat_hitl.environment.controllers.llm_controller import (
+    AgentTerminationEvent,
     LLMController,
     PlannerStatus,
 )
@@ -57,7 +58,6 @@ class EpisodeCompletionStatus(Enum):
     PENDING = (0,)
     SUCCESS = (1,)
     FAILURE = (2,)
-    ERROR   = (3,)
 
 
 class FrameRecorder:
@@ -166,15 +166,15 @@ def __init__(
         self.cam_transform = mn.Matrix4.identity_init()
 
         self.episode_completion_status = EpisodeCompletionStatus.PENDING
-        self._episode_completion_message = ''
+        self._episode_completion_message = ""
 
-    def _on_termination_cb(self, _e: Any = None):
+    def _on_termination_cb(self, _e: AgentTerminationEvent = None):
         if _e.status == PlannerStatus.SUCCESS:
             self.episode_completion_status = EpisodeCompletionStatus.SUCCESS
         elif _e.status == PlannerStatus.FAILED:
             self.episode_completion_status = EpisodeCompletionStatus.FAILURE
         elif _e.status == PlannerStatus.ERROR:
-            self.episode_completion_status = EpisodeCompletionStatus.ERROR
+            self.episode_completion_status = EpisodeCompletionStatus.FAILURE
             self._episode_message = _e.message
 
     def update_camera_from_sensor(self) -> None:
@@ -639,7 +639,9 @@ def on_exit(self):
         super().on_exit()
 
         episode_success = self._is_episode_successful()
-        self._session.session_recorder.end_episode(episode_success)
+        self._session.session_recorder.end_episode(
+            episode_success, self._skip_episode_error_message
+        )
 
     def on_environment_reset(self, episode_recorder_dict):
         self._world.reset()
@@ -920,7 +922,7 @@ def _is_episode_successful(self) -> bool:
             for agent_index in range(self._num_agents)
         )
 
-    def _on_termination_cb(self, _e: Any = None):
+    def _on_termination_cb(self, _e: AgentTerminationEvent = None):
         # Trigger episode change sequence when an agent error occurs.
         if _e.status == PlannerStatus.ERROR:
             self._skip_episode_error_message = f"Other participant encountered an error: {_e.message}. Skipping episode."

examples/hitl/rearrange_v2/session_recorder.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 from util import timestamp
 
@@ -68,10 +68,14 @@ def start_episode(
     def end_episode(
         self,
         success: bool,
+        episode_failure_error_message: Optional[str],
     ):
         self.data["episodes"][-1]["end_timestamp"] = timestamp()
         self.data["episodes"][-1]["success"] = success
         self.data["episodes"][-1]["completed"] = True
+        self.data["episodes"][-1][
+            "error_message"
+        ] = episode_failure_error_message
 
     def record_frame(
         self,

habitat-hitl/habitat_hitl/environment/controllers/llm_controller.py

@@ -270,15 +270,23 @@ def act(self, observations, debug_obs: bool = False, *args, **kwargs):
                 >= self._planner_info["replanning_threshold"]
             ):
                 self._on_termination.invoke(
-                    AgentTerminationEvent(status=PlannerStatus.FAILED, message="replanning threshold exceeded")
+                    AgentTerminationEvent(
+                        status=PlannerStatus.FAILED,
+                        message="replanning threshold exceeded",
+                    )
                 )
             elif "ConnectionError" in self._planner_info["prompts"][0]:
                 self._on_termination.invoke(
-                    AgentTerminationEvent(status=PlannerStatus.ERROR, message="LLM connection error")
+                    AgentTerminationEvent(
+                        status=PlannerStatus.ERROR,
+                        message="LLM connection error",
+                    )
                 )
             else:
                 self._on_termination.invoke(
-                    AgentTerminationEvent(status=PlannerStatus.SUCCESS, message="")
+                    AgentTerminationEvent(
+                        status=PlannerStatus.SUCCESS, message=""
+                    )
                 )
             self._termination_reported = True