Loop around level loader on finishing all levels (#13)

Author: taufeeque9

envpool/sokoban/level_loader.cc

@@ -32,6 +32,7 @@ LevelLoader::LevelLoader(const std::filesystem::path& base_path,
                          int env_id, int num_envs, int verbose)
     : load_sequentially_(load_sequentially),
       n_levels_to_load_(n_levels_to_load),
+      env_id_(env_id),
       num_envs_(num_envs),
       cur_level_(env_id),
       verbose(verbose) {
@@ -54,6 +55,7 @@ LevelLoader::LevelLoader(const std::filesystem::path& base_path,
     throw std::runtime_error(
         "n_levels_to_load must be a multiple of num_envs.");
   }
+  n_levels_to_load_ /= num_envs_;
 }
 
 static const std::array<char, kMaxLevelObject + 1> kPrintLevelKey{
@@ -186,16 +188,19 @@ void LevelLoader::LoadFile(std::mt19937& gen) {
 
 std::vector<SokobanLevel>::iterator LevelLoader::GetLevel(std::mt19937& gen) {
   if (n_levels_to_load_ > 0 && levels_loaded_ >= n_levels_to_load_) {
-    throw std::runtime_error("Loaded all requested levels.");
+    // std::cerr << "Warning: All levels loaded. Looping around now." <<
+    // std::endl;
+    levels_loaded_ = 0;
+    cur_file_ = level_file_paths_.begin();
+    LoadFile(gen);
+    // re-start from the `env_id`th level, like we do in the constructor.
+    cur_level_ = env_id_;
   }
   // Load new files until the current level index is within the loaded levels
   // this is required when new files have lesser levels than the number of envs
   while (cur_level_ >= levels_.size()) {
     cur_level_ -= levels_.size();
     LoadFile(gen);
-    if (levels_.empty()) {  // new file is empty
-      throw std::runtime_error("No levels loaded.");
-    }
   }
   // no need for bound checks since it is checked in the while loop above
   auto out = levels_.begin() + cur_level_;

envpool/sokoban/sokoban_py_envpool_test.py

@@ -305,19 +305,21 @@ def test_load_sequentially_with_multiple_envs() -> None:
   )
   dim_room = env.spec.config.dim_room
   printed_obs = []
-  for _ in range(total_levels // num_envs):
-    obs, _ = env.reset()
-    assert obs.shape == (
-      num_envs,
-      3,
-      dim_room,
-      dim_room,
-    ), f"obs shape: {obs.shape}"
-    for idx in range(num_envs):
-      printed_obs.append(print_obs(obs[idx]).strip().split("\n"))
-  for i, level in enumerate(levels_by_files):
-    for j, line in enumerate(level):
-      assert printed_obs[i][j] == line, f"Level {i} is not loaded correctly."
+
+  for _ in range(2):  # check loader loops around and loads levels again
+    for _ in range(total_levels // num_envs):
+      obs, _ = env.reset()
+      assert obs.shape == (
+        num_envs,
+        3,
+        dim_room,
+        dim_room,
+      ), f"obs shape: {obs.shape}"
+      for idx in range(num_envs):
+        printed_obs.append(print_obs(obs[idx]).strip().split("\n"))
+    for i, level in enumerate(levels_by_files):
+      for j, line in enumerate(level):
+        assert printed_obs[i][j] == line, f"Level {i} is not loaded correctly."
 
 
 def test_astar_log(tmp_path) -> None: