Episode iterator upgrades

configs/test/habitat_all_sensors_test.yaml

@@ -1,5 +1,7 @@
 ENVIRONMENT:
   MAX_EPISODE_STEPS: 10
+  ITERATOR_OPTIONS:
+    SHUFFLE: False
 SIMULATOR:
   AGENT_0:
     SENSORS: ['RGB_SENSOR', 'DEPTH_SENSOR']

configs/test/habitat_mp3d_eqa_test.yaml

@@ -1,5 +1,8 @@
 TASK:
   TYPE: EQA-v0
+ENVIRONMENT:
+  ITERATOR_OPTIONS:
+    SHUFFLE: False
 SIMULATOR:
   SCENE: data/scene_datasets/mp3d/17DRP5sb8fy/17DRP5sb8fy.glb
   FORWARD_STEP_SIZE: 0.1

habitat/config/default.py

@@ -24,10 +24,11 @@
 _C.ENVIRONMENT.MAX_EPISODE_SECONDS = 10000000
 _C.ENVIRONMENT.ITERATOR_OPTIONS = CN()
 _C.ENVIRONMENT.ITERATOR_OPTIONS.CYCLE = True
-_C.ENVIRONMENT.ITERATOR_OPTIONS.SHUFFLE = False
+_C.ENVIRONMENT.ITERATOR_OPTIONS.SHUFFLE = True
 _C.ENVIRONMENT.ITERATOR_OPTIONS.GROUP_BY_SCENE = True
 _C.ENVIRONMENT.ITERATOR_OPTIONS.NUM_EPISODE_SAMPLE = -1
-_C.ENVIRONMENT.ITERATOR_OPTIONS.MAX_SCENE_REPEAT = -1
+_C.ENVIRONMENT.ITERATOR_OPTIONS.MAX_SCENE_REPEAT_EPISODES = -1
+_C.ENVIRONMENT.ITERATOR_OPTIONS.MAX_SCENE_REPEAT_STEPS = int(1e4)
 # -----------------------------------------------------------------------------
 # TASK
 # -----------------------------------------------------------------------------

habitat/core/dataset.py

@@ -278,7 +278,8 @@ def __init__(
         cycle: bool = True,
         shuffle: bool = False,
         group_by_scene: bool = True,
-        max_scene_repeat: int = -1,
+        max_scene_repeat_episodes: int = -1,
+        max_scene_repeat_steps: int = -1,
         num_episode_sample: int = -1,
     ):
         r"""..
@@ -290,32 +291,43 @@ def __init__(
             effect if cycle is set to :py:`False`. Will shuffle grouped scenes
             if :p:`group_by_scene` is :py:`True`.
         :param group_by_scene: if :py:`True`, group episodes from same scene.
-        :param max_scene_repeat: threshold of how many episodes from the same
+        :param max_scene_repeat_episodes: threshold of how many episodes from the same
             scene can be loaded consecutively. :py:`-1` for no limit
+        :param max_scene_repeat_steps: threshold of how many steps from the same
+            scene can be taken consecutively. :py:`-1` for no limit
         :param num_episode_sample: number of episodes to be sampled. :py:`-1`
             for no sampling.
         """
+        self._repetition_rand_interval = 0.2
+
         # sample episodes
         if num_episode_sample >= 0:
             episodes = np.random.choice(
                 episodes, num_episode_sample, replace=False
             )
+
         self.episodes = episodes
         self.cycle = cycle
         self.group_by_scene = group_by_scene
-        if group_by_scene:
-            num_scene_groups = len(
-                list(groupby(episodes, key=lambda x: x.scene_id))
-            )
-            num_unique_scenes = len(set([e.scene_id for e in episodes]))
-            if num_scene_groups >= num_unique_scenes:
-                self.episodes = sorted(self.episodes, key=lambda x: x.scene_id)
-        self.max_scene_repetition = max_scene_repeat
         self.shuffle = shuffle
+
+        if shuffle:
+            random.shuffle(self.episodes)
+
+        if group_by_scene:
+            self.episodes = sorted(self.episodes, key=lambda x: x.scene_id)
+
+        self.max_scene_repetition_episodes = max_scene_repeat_episodes
+        self.max_scene_repetition_steps = max_scene_repeat_steps
+
         self._rep_count = 0
+        self._step_count = 0
         self._prev_scene_id = None
+
         self._iterator = iter(self.episodes)
 
+        self._set_shuffle_intervals()
+
     def __iter__(self):
         return self
 
@@ -324,6 +336,7 @@ def __next__(self):
 
         :return: next episode.
         """
+        self._switch_scene_if()
 
         next_episode = next(self._iterator, None)
         if next_episode is None:
@@ -334,14 +347,9 @@ def __next__(self):
                 self._shuffle_iterator()
             next_episode = next(self._iterator)
 
-        if self._prev_scene_id == next_episode.scene_id:
-            self._rep_count += 1
-        if (
-            self.max_scene_repetition > 0
-            and self._rep_count >= self.max_scene_repetition - 1
-        ):
-            self._shuffle_iterator()
+        if self._prev_scene_id != next_episode.scene_id:
             self._rep_count = 0
+            self._step_count = 0
 
         self._prev_scene_id = next_episode.scene_id
         return next_episode
@@ -355,9 +363,60 @@ def _shuffle_iterator(self) -> None:
                 list(g)
                 for k, g in groupby(self._iterator, key=lambda x: x.scene_id)
             ]
+
             random.shuffle(grouped_episodes)
+            for i in range(len(grouped_episodes)):
+                random.shuffle(grouped_episodes[i])
+
             self._iterator = iter(sum(grouped_episodes, []))
         else:
             episodes = list(self._iterator)
             random.shuffle(episodes)
             self._iterator = iter(episodes)
+
+    def step_taken(self):
+        self._step_count += 1
+
+    @staticmethod
+    def _randomize_value(value, interval):
+        return random.randint(
+            int(value * (1 - interval)), int(value * (1 + interval))
+        )
+
+    def _set_shuffle_intervals(self):
+        if self.max_scene_repetition_episodes > 0:
+            self._max_rep_episode = self._randomize_value(
+                self.max_scene_repetition_episodes,
+                self._repetition_rand_interval,
+            )
+        else:
+            self._max_rep_episode = None
+
+        if self.max_scene_repetition_steps > 0:
+            self._max_rep_step = self._randomize_value(
+                self.max_scene_repetition_steps, self._repetition_rand_interval
+            )
+        else:
+            self._max_rep_step = None
+
+    def _switch_scene_if(self):
+        do_switch = False
+        self._rep_count += 1
+
+        # Shuffle if a scene has been selected more than _max_rep_episode times in a row
+        if (
+            self._max_rep_episode is not None
+            and self._rep_count > self._max_rep_episode
+        ):
+            do_switch = True
+
+        # Shuffle if a scene has been used for more than _max_rep_step steps in a row
+        if (
+            self._max_rep_step is not None
+            and self._step_count >= self._max_rep_step
+        ):
+            do_switch = True
+
+        if do_switch:
+            self._shuffle_iterator()
+            self._set_shuffle_intervals()

habitat/core/env.py

@@ -12,7 +12,7 @@
 from gym.spaces.dict_space import Dict as SpaceDict
 
 from habitat.config import Config
-from habitat.core.dataset import Dataset, Episode
+from habitat.core.dataset import Dataset, Episode, EpisodeIterator
 from habitat.core.embodied_task import EmbodiedTask, Metrics
 from habitat.core.simulator import Observations, Simulator
 from habitat.datasets import make_dataset
@@ -213,6 +213,11 @@ def _update_step_stats(self) -> None:
         if self._past_limit():
             self._episode_over = True
 
+        if self.episode_iterator is not None and isinstance(
+            self.episode_iterator, EpisodeIterator
+        ):
+            self.episode_iterator.step_taken()
+
     def step(
         self, action: Union[int, str, Dict[str, Any]], **kwargs
     ) -> Observations:

test/test_dataset.py

@@ -240,14 +240,20 @@ def test_iterator_shuffle():
     assert len(first_round_scene_groups) == len(set(first_round_scene_groups))
 
 
-def test_iterator_scene_switching():
+def test_iterator_scene_switching_episodes():
     total_ep = 1000
     max_repeat = 25
-    dataset = _construct_dataset(total_ep)
+    dataset = _construct_dataset(total_ep, num_groups=5)
 
-    episode_iter = dataset.get_episode_iterator(max_scene_repeat=max_repeat)
+    episode_iter = dataset.get_episode_iterator(
+        max_scene_repeat_episodes=max_repeat, shuffle=False
+    )
+    episode_iter._repetition_rand_interval = 0
+    episode_iter._set_shuffle_intervals()
     episodes = sorted(dataset.episodes, key=lambda x: x.scene_id)
 
+    max_repeat = episode_iter._max_rep_episode
+
     # episodes before max_repeat reached should be identical
     for i in range(max_repeat):
         episode = next(episode_iter)
@@ -260,4 +266,39 @@ def test_iterator_scene_switching():
     assert sorted(remaining_episodes) == sorted(episodes)
 
     # next episodes should still be grouped by scene (before next switching)
-    assert len(set([e.scene_id for e in remaining_episodes[:max_repeat]])) == 1
+    assert len(set(e.scene_id for e in remaining_episodes)) == len(
+        list(groupby(remaining_episodes, lambda ep: ep.scene_id))
+    )
+
+
+def test_iterator_scene_switching_steps():
+    total_ep = 1000
+    max_repeat_steps = 250
+    dataset = _construct_dataset(total_ep, num_groups=10)
+
+    episode_iter = dataset.get_episode_iterator(
+        max_scene_repeat_steps=max_repeat_steps, shuffle=False
+    )
+    episode_iter._repetition_rand_interval = 0
+    episode_iter._set_shuffle_intervals()
+    episodes = sorted(dataset.episodes, key=lambda x: x.scene_id)
+
+    max_repeat_steps = episode_iter._max_rep_step
+
+    episode = next(episode_iter)
+    assert episode.episode_id == episodes.pop(0).episode_id
+
+    # episodes before max_repeat reached should be identical
+    for _ in range(max_repeat_steps):
+        episode_iter.step_taken()
+
+    episode = next(episode_iter)
+    assert episode.episode_id != episodes.pop(0).episode_id
+
+    remaining_episodes = list(islice(episode_iter, total_ep - 2))
+    assert len(remaining_episodes) == len(episodes)
+
+    # next episodes should still be grouped by scene (before next switching)
+    assert len(set(e.scene_id for e in remaining_episodes)) == len(
+        list(groupby(remaining_episodes, lambda ep: ep.scene_id))
+    )

test/test_habitat_task.py

@@ -35,18 +35,12 @@ def test_task_actions():
         }
     )
     agent_state = env.sim.get_agent_state()
-    assert (
-        np.allclose(
-            np.array(TELEPORT_POSITION, dtype=np.float32), agent_state.position
-        )
-        is True
+    assert np.allclose(
+        np.array(TELEPORT_POSITION, dtype=np.float32), agent_state.position
     ), "mismatch in position after teleport"
-    assert (
-        np.allclose(
-            np.array(TELEPORT_ROTATION, dtype=np.float32),
-            np.array([*agent_state.rotation.imag, agent_state.rotation.real]),
-        )
-        is True
+    assert np.allclose(
+        np.array(TELEPORT_ROTATION, dtype=np.float32),
+        np.array([*agent_state.rotation.imag, agent_state.rotation.real]),
     ), "mismatch in rotation after teleport"
     env.step("TURN_RIGHT")
     env.close()