[Bugfix] [Rearrange] Reset the sensors positons correctly version 2

habitat/tasks/rearrange/multi_task/composite_task.py

@@ -66,4 +66,5 @@ def reset(self, episode: Episode):
         if self._cur_node_idx >= 0:
             self.jump_to_node(self._cur_node_idx, episode)
 
+        self._sim.maybe_update_robot()
         return self._get_observations(episode)

habitat/tasks/rearrange/rearrange_sim.py

@@ -628,8 +628,7 @@ def step(self, action: Union[str, int]) -> Observations:
                 add_back_viz_objs[name] = (before_pos, r)
             self.viz_ids = defaultdict(lambda: None)
 
-        if self.habitat_config.UPDATE_ROBOT:
-            self.robots_mgr.update_robots()
+        self.maybe_update_robot()
 
         if self.habitat_config.CONCUR_RENDER:
             self._prev_sim_obs = self.start_async_render()
@@ -670,6 +669,16 @@ def step(self, action: Union[str, int]) -> Observations:
 
         return obs
 
+    def maybe_update_robot(self):
+        """
+        Calls the update robots method on the robot manager if the
+        `UPDATE_ROBOT` configuration is set to True. Among other
+        things, this will set the robot's sensors' positions to their new
+        positions.
+        """
+        if self.habitat_config.UPDATE_ROBOT:
+            self.robots_mgr.update_robots()
+
     def visualize_position(
         self,
         position: np.ndarray,

habitat/tasks/rearrange/rearrange_task.py

@@ -168,6 +168,7 @@ def reset(self, episode: Episode, fetch_observations: bool = True):
         self._done = False
         self._cur_episode_step = 0
         if fetch_observations:
+            self._sim.maybe_update_robot()
             return self._get_observations(episode)
         else:
             return None

habitat/tasks/rearrange/sub_tasks/articulated_object_task.py

@@ -186,6 +186,7 @@ def reset(self, episode: Episode):
         self.prev_dist_to_push = -1
 
         self.prev_snapped_marker_name = None
+        self._sim.maybe_update_robot()
         return self._get_observations(episode)
 
     def _disable_art_sleep(self):

habitat/tasks/rearrange/sub_tasks/nav_to_obj_task.py

@@ -343,7 +343,7 @@ def reset(self, episode: Episode):
                 sim.viz_ids["nav_targ_pos"],
                 r=0.2,
             )
-
+        self._sim.maybe_update_robot()
         return self._get_observations(episode)
 
 

habitat/tasks/rearrange/sub_tasks/pick_task.py

@@ -283,5 +283,6 @@ def reset(self, episode: Episode, fetch_observations: bool = True):
         self._targ_idx = sel_idx
 
         if fetch_observations:
+            self._sim.maybe_update_robot()
             return self._get_observations(episode)
         return None

habitat/tasks/rearrange/sub_tasks/place_task.py

@@ -41,4 +41,5 @@ def reset(self, episode: Episode):
         self.was_prev_holding = self.targ_idx
 
         sim.internal_step(-1)
+        self._sim.maybe_update_robot()
         return self._get_observations(episode)

habitat/tasks/rearrange/sub_tasks/reach_task.py

@@ -61,4 +61,5 @@ def reset(self, episode):
                 global_pos, self._sim.viz_ids["reach_target"]
             )
 
+        self._sim.maybe_update_robot()
         return self._get_observations(episode)