reset to demo bug fix

stepjam · Jun 11, 2024 · c7c50ed · c7c50ed
1 parent 6e97fef
commit c7c50ed
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/rlbench/backend/scene.py b/rlbench/backend/scene.py
@@ -134,7 +134,10 @@ def init_episode(self, index: int, randomly_place: bool=True,
                         raise BoundaryError()
                 if not place_demo:
                     self.task.validate()
-                break
+                    break
+                else:
+                    # Placing demo, run the number of attempts for correct demo reset
+                    self._attempts += 1
             except (BoundaryError, WaypointError) as e:
                 self.task.cleanup_()
                 self.task.restore_state(self._initial_task_state)

diff --git a/rlbench/task_environment.py b/rlbench/task_environment.py
@@ -77,7 +77,7 @@ def reset(self, demo = None) -> (List[str], Observation):
         try:
             place_demo = demo != None and hasattr(demo, 'num_reset_attempts') and demo.num_reset_attempts != None
             desc = self._scene.init_episode(
-                self._variation_number, max_attempts=_MAX_RESET_ATTEMPTS,
+                self._variation_number, max_attempts=_MAX_RESET_ATTEMPTS if not place_demo else demo.num_reset_attempts,
                 randomly_place=not self._static_positions, place_demo=place_demo)
         except (BoundaryError, WaypointError) as e:
             raise TaskEnvironmentError(