ignoring out of bounds error and calling stop instead

bdaiinstitute · Sep 2, 2023 · fd10075 · fd10075
1 parent ee986fe
commit fd10075
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/zsos/policy/base_objectnav_policy.py b/zsos/policy/base_objectnav_policy.py
@@ -143,7 +143,12 @@ def _pre_step(self, observations: "TensorDict", masks: Tensor) -> None:
         if not self._did_reset and masks[0] == 0:
             self._reset()
             self._target_object = observations["objectgoal"]
-        self._cache_observations(observations)
+        try:
+            self._cache_observations(observations)
+        except IndexError as e:
+            print(e)
+            print("Reached edge of map, stopping.")
+            raise StopIteration
         self._policy_info = {}
 
     def _initialize(self) -> Tensor:

diff --git a/zsos/policy/habitat_policies.py b/zsos/policy/habitat_policies.py
@@ -133,9 +133,12 @@ def act(
         else:
             raise ValueError(f"Dataset type {self._dataset_type} not recognized")
         parent_cls: BaseObjectNavPolicy = super()  # type: ignore
-        action, rnn_hidden_states = parent_cls.act(
-            obs_dict, rnn_hidden_states, prev_actions, masks, deterministic
-        )
+        try:
+            action, rnn_hidden_states = parent_cls.act(
+                obs_dict, rnn_hidden_states, prev_actions, masks, deterministic
+            )
+        except StopIteration:
+            action = self._stop_action
         return PolicyActionData(
             actions=action,
             rnn_hidden_states=rnn_hidden_states,