Add option to load a fixed number of levels sequentially

envpool/sokoban/level_loader.cc

@@ -1,5 +1,7 @@
 #include "level_loader.h"
 
+#include <pybind11/iostream.h>
+
 #include <algorithm>
 #include <filesystem>
 #include <fstream>
@@ -12,14 +14,23 @@ namespace sokoban {
 
 size_t ERROR_SZ = 1024;
 
-LevelLoader::LevelLoader(const std::filesystem::path& base_path, int verbose)
-    : levels(0), cur_level(levels.begin()), level_file_paths(0), verbose(verbose) {
+LevelLoader::LevelLoader(const std::filesystem::path& base_path,
+                         bool load_sequentially, int n_levels_to_load,
+                         int verbose)
+    : load_sequentially(load_sequentially),
+      n_levels_to_load(n_levels_to_load),
+      levels_loaded(0),
+      levels(0),
+      cur_level(levels.begin()),
+      level_file_paths(0),
+      verbose(verbose) {
   for (const auto& entry : std::filesystem::directory_iterator(base_path)) {
     level_file_paths.push_back(entry.path());
   }
+  cur_file = level_file_paths.begin();
 }
 
-const std::string PRINT_LEVEL_KEY = "# .a@$s";
+const std::string PRINT_LEVEL_KEY = "# .a$@s";
 
 void AddLine(SokobanLevel& level, const std::string& line) {
   auto start = line.at(0);
@@ -69,11 +80,20 @@ void PrintLevel(std::ostream& os, SokobanLevel vec) {
   }
 }
 
-void LevelLoader::LoadNewFile(std::mt19937& gen) {
-  std::uniform_int_distribution<size_t> load_file_idx_r(
-      0, level_file_paths.size() - 1);
-  const size_t load_file_idx = load_file_idx_r(gen);
-  const std::filesystem::path& file_path = level_file_paths.at(load_file_idx);
+void LevelLoader::LoadFile(std::mt19937& gen) {
+  std::filesystem::path file_path;
+  if (load_sequentially) {
+    if (cur_file == level_file_paths.end()) {
+      throw std::runtime_error("No more files to load.");
+    }
+    file_path = *cur_file;
+    cur_file++;
+  } else {
+    std::uniform_int_distribution<size_t> load_file_idx_r(
+        0, level_file_paths.size() - 1);
+    const size_t load_file_idx = load_file_idx_r(gen);
+    file_path = level_file_paths.at(load_file_idx);
+  }
   std::ifstream file(file_path);
 
   levels.clear();
@@ -114,17 +134,19 @@ void LevelLoader::LoadNewFile(std::mt19937& gen) {
       }
     }
   }
-  std::shuffle(levels.begin(), levels.end(), gen);
+  if (!load_sequentially) {
+    std::shuffle(levels.begin(), levels.end(), gen);
+  }
   if (levels.empty()) {
     std::stringstream msg;
     msg << "No levels loaded from file '" << file_path << std::endl;
     throw std::runtime_error(msg.str());
   }
 
-  if(verbose >= 1) {
-    std::cout << "Loaded " << levels.size() << " levels from " << file_path
+  if (verbose >= 1) {
+    std::cout << "***Loaded " << levels.size() << " levels from " << file_path
               << std::endl;
-    if(verbose >= 2) {
+    if (verbose >= 2) {
       PrintLevel(std::cout, levels.at(0));
       std::cout << std::endl;
       PrintLevel(std::cout, levels.at(1));
@@ -133,17 +155,21 @@ void LevelLoader::LoadNewFile(std::mt19937& gen) {
   }
 }
 
-const std::vector<SokobanLevel>::iterator LevelLoader::RandomLevel(
+const 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.");
+  }
   if (cur_level == levels.end()) {
-    LoadNewFile(gen);
+    LoadFile(gen);
     cur_level = levels.begin();
     if (cur_level == levels.end()) {
       throw std::runtime_error("No levels loaded.");
     }
   }
   auto out = cur_level;
   cur_level++;
+  levels_loaded++;
   return out;
 }
 

envpool/sokoban/level_loader.h

@@ -19,19 +19,23 @@ constexpr uint8_t PLAYER_ON_TARGET = 6;
 
 class LevelLoader {
  protected:
+  bool load_sequentially;
+  int n_levels_to_load;
+  int levels_loaded;
   std::vector<SokobanLevel> levels;
   std::vector<SokobanLevel>::iterator cur_level;
   std::vector<std::filesystem::path> level_file_paths;
-  void LoadNewFile(std::mt19937& gen);
+  std::vector<std::filesystem::path>::iterator cur_file;
+  void LoadFile(std::mt19937& gen);
 
  public:
   int verbose;
 
-  const std::vector<SokobanLevel>::iterator RandomLevel(std::mt19937& gen);
-  LevelLoader(const std::filesystem::path& base_path, int verbose=0);
+  const std::vector<SokobanLevel>::iterator GetLevel(std::mt19937& gen);
+  LevelLoader(const std::filesystem::path& base_path, bool load_sequentially,
+              int n_levels_to_load, int verbose = 0);
 };
 
-
 void PrintLevel(std::ostream& os, SokobanLevel vec);
 }  // namespace sokoban
 

envpool/sokoban/registration.py

@@ -9,4 +9,5 @@
     gymnasium_cls="SokobanGymnasiumEnvPool",
     max_episode_steps=60,
     reward_step=-0.1,
+    max_num_players=1,
 )

envpool/sokoban/sample_levels/001.txt

@@ -0,0 +1,35 @@
+; 0
+##########
+##########
+##########
+##### # ##
+#####    #
+#####  $ #
+# .    ..#
+# $$$ #  #
+#@ .     #
+##########
+
+; 1
+##########
+##########
+####     #
+#  $  .  #
+# #      #
+#@### .$ #
+###### $ #
+###  $.  #
+###     .#
+##########
+
+; 2
+##########
+#####   @#
+####    ##
+####.   ##
+#.   . $ #
+# $  $.  #
+# ###    #
+# $ ######
+#   ######
+##########

envpool/sokoban/sokoban_envpool.cc

@@ -11,10 +11,11 @@ namespace sokoban {
 void SokobanEnv::Reset() {
   const int max_episode_steps = spec_.config["max_episode_steps"_];
   const int min_episode_steps = spec_.config["min_episode_steps"_];
-  std::uniform_int_distribution<int> episode_length_rand(min_episode_steps, max_episode_steps);
+  std::uniform_int_distribution<int> episode_length_rand(min_episode_steps,
+                                                         max_episode_steps);
   current_max_episode_steps_ = episode_length_rand(gen_);
 
-  world = *level_loader.RandomLevel(gen_);
+  world = *level_loader.GetLevel(gen_);
   if (world.size() != dim_room * dim_room) {
     std::stringstream msg;
     msg << "Loaded level is not dim_room x dim_room. world.size()="
@@ -55,30 +56,17 @@ void SokobanEnv::WorldAssignAt(int x, int y, uint8_t value) {
 constexpr std::array<std::array<int, 2>, 4> CHANGE_COORDINATES = {
     {{0, -1}, {0, 1}, {-1, 0}, {1, 0}}};
 
-constexpr std::array<const char *, MAX_ACTION+1> action_names = {
-  "ACT_NOOP",
-  "ACT_PUSH_UP",
-  "ACT_PUSH_DOWN",
-  "ACT_PUSH_LEFT",
-  "ACT_PUSH_RIGHT",
-  "ACT_MOVE_UP",
-  "ACT_MOVE_DOWN",
-  "ACT_MOVE_LEFT",
-  "ACT_MOVE_RIGHT",
+constexpr std::array<const char*, MAX_ACTION + 1> action_names = {
+    "ACT_NOOP",      "ACT_PUSH_UP",    "ACT_PUSH_DOWN",
+    "ACT_PUSH_LEFT", "ACT_PUSH_RIGHT", "ACT_MOVE_UP",
+    "ACT_MOVE_DOWN", "ACT_MOVE_LEFT",  "ACT_MOVE_RIGHT",
 };
 
-
-constexpr std::array<const char *, 7> arena_names = {
-  "WALL",
-  "EMPTY",
-  "TARGET",
-  "BOX_ON_TARGET",
-  "BOX",
-  "PLAYER",
-  "PLAYER_ON_TARGET",
+constexpr std::array<const char*, 7> arena_names = {
+    "WALL", "EMPTY",  "TARGET",           "BOX_ON_TARGET",
+    "BOX",  "PLAYER", "PLAYER_ON_TARGET",
 };
 
-
 void SokobanEnv::Step(const Action& action_) {
   current_step_++;
 

envpool/sokoban/sokoban_envpool.h

@@ -29,7 +29,9 @@ class SokobanEnvFns {
     return MakeDict("reward_finished"_.Bind(10.0), "reward_box"_.Bind(1.0),
                     "reward_step"_.Bind(-0.1), "dim_room"_.Bind(10),
                     "levels_dir"_.Bind(std::string("")), "verbose"_.Bind(0),
-                    "min_episode_steps"_.Bind(0));
+                    "min_episode_steps"_.Bind(0),
+                    "load_sequentially"_.Bind(false),
+                    "n_levels_to_load"_.Bind(-1));
   }
   template <typename Config>
   static decltype(auto) StateSpec(const Config& conf) {
@@ -54,10 +56,13 @@ class SokobanEnv : public Env<SokobanEnvSpec> {
         reward_box{static_cast<double>(spec.config["reward_box"_])},
         reward_step{static_cast<double>(spec.config["reward_step"_])},
         levels_dir{static_cast<std::string>(spec.config["levels_dir"_])},
-        level_loader(levels_dir),
+        level_loader(levels_dir, spec.config["load_sequentially"_],
+                     static_cast<int>(spec.config["n_levels_to_load"_]),
+                     static_cast<int>(spec.config["verbose"_])),
         world(WALL, static_cast<std::size_t>(dim_room * dim_room)),
         verbose(static_cast<int>(spec.config["verbose"_])),
-        current_max_episode_steps_(static_cast<int>(spec.config["max_episode_steps"_])) {
+        current_max_episode_steps_(
+            static_cast<int>(spec.config["max_episode_steps"_])) {
     if (max_num_players_ != spec_.config["max_num_players"_]) {
       std::stringstream msg;
       msg << "max_num_players_ != spec_['max_num_players'] " << max_num_players_
@@ -74,7 +79,9 @@ class SokobanEnv : public Env<SokobanEnvSpec> {
   }
 
   bool IsDone() override {
-    return (unmatched_boxes == 0) || (current_step_ >= current_max_episode_steps_); }
+    return (unmatched_boxes == 0) ||
+           (current_step_ >= current_max_episode_steps_);
+  }
   void Reset() override;
   void Step(const Action& action) override;
 

envpool/sokoban/sokoban_py_envpool_test.py

@@ -1,6 +1,9 @@
 """Unit test for dummy envpool and speed benchmark."""
 
+import glob
 import time
+import py
+import re
 
 import envpool  # noqa: F401
 import envpool.sokoban.registration
@@ -32,6 +35,8 @@ def test_config(self) -> None:
             "reward_finished",
             "reward_step",
             "verbose",
+            "load_sequentially",
+            "n_levels_to_load",
         ]
         default_conf = _SokobanEnvSpec._default_config_values
         self.assertTrue(isinstance(default_conf, tuple))
@@ -71,6 +76,7 @@ def test_envpool_max_episode_steps(self) -> None:
                 num_envs=1,
                 batch_size=1,
                 max_episode_steps=max_episode_steps,
+                min_episode_steps=max_episode_steps,
                 levels_dir="/app/envpool/sokoban/sample_levels",
             )
             env.reset()
@@ -82,6 +88,51 @@ def test_envpool_max_episode_steps(self) -> None:
             assert not np.any(terminated)
             assert np.all(truncated)
 
+    def test_envpool_load_sequentially(self) -> None:
+        levels_dir = "/app/envpool/sokoban/sample_levels"
+        files = glob.glob(f"{levels_dir}/*.txt")
+        levels_by_files = []
+        for file in files:
+            with open(file, "r") as f:
+                text = f.read()
+            levels = text.split("\n;")
+            levels = ["\n".join(level.split("\n")[1:]).strip() for level in levels]
+            levels_by_files.append((file, levels))
+        assert len(levels_by_files) > 1
+        assert all(len(levels) > 1 for levels in levels_by_files)
+        total_levels = sum(len(levels) for levels in levels_by_files)
+        for n_levels_to_load in range(1, total_levels + 1):
+            env = envpool.make(
+                "Sokoban-v0",
+                env_type="gymnasium",
+                num_envs=1,
+                batch_size=1,
+                max_episode_steps=60,
+                min_episode_steps=60,
+                levels_dir=levels_dir,
+                load_sequentially=True,
+                n_levels_to_load=n_levels_to_load,
+                verbose=2,
+            )
+            dim_room = env.spec.config.dim_room
+            capture = py.io.StdCaptureFD()
+            obs, _ = env.reset()
+            assert obs.shape == (1, 3, dim_room, dim_room), f"obs shape: {obs.shape}"
+            if n_levels_to_load == -1:
+                n_levels_to_load = total_levels
+            for _ in range(n_levels_to_load - 1):
+                env.reset()
+            out, err = capture.reset()
+            files_output = out.split("***")[1:]
+            for i, file_output in enumerate(files_output):
+                first_line, out = file_output.strip().split("\n", 1)
+                result = re.search(r'Loaded (\d+) levels from "(.*\.txt)"', first_line)
+                n_levels, file_name = int(result.group(1)), result.group(2)
+                lev1, lev2 = out.strip().split("\n\n")
+                assert file_name == levels_by_files[i][0]
+                assert lev1 == levels_by_files[i][1][0]
+                assert lev2 == levels_by_files[i][1][1]
+
     def test_xla(self) -> None:
         num_envs = 10
         env = envpool.make(