Add option to load a fixed number of levels sequentially

envpool/sokoban/BUILD

@@ -15,30 +15,31 @@
 load("@pip_requirements//:requirements.bzl", "requirement")
 load("@pybind11_bazel//:build_defs.bzl", "pybind_extension")
 
-package(default_visibility = ["//visibility:public"])
+package(default_visibility=["//visibility:public"])
 
 py_library(
-    name = "sokoban",
-    srcs = ["__init__.py"],
-    data = [":sokoban_envpool.so"],
-    deps = ["//envpool/python:api"],
+    name="sokoban",
+    srcs=["__init__.py"],
+    data=[":sokoban_envpool.so"],
+    deps=["//envpool/python:api"],
 )
 
 py_library(
-    name = "registration",
-    srcs = ["registration.py"],
-    deps = [
+    name="registration",
+    srcs=["registration.py"],
+    deps=[
         "//envpool:registration",
     ],
 )
 
 cc_library(
-    name = "sokoban_envpool_h",
-    hdrs = [
+    name="sokoban_envpool_h",
+    hdrs=[
         "level_loader.h",
         "sokoban_envpool.h",
+        "utils.h",
     ],
-    deps = [
+    deps=[
         "//envpool/core:async_envpool",
         "//envpool/core:env",
         "//envpool/core:env_spec",
@@ -56,28 +57,28 @@ cc_library(
 # )
 
 py_test(
-    name = "test",
-    srcs = ["sokoban_py_envpool_test.py"],
-    main = "sokoban_py_envpool_test.py",
-    deps = [
+    name="test",
+    srcs=["sokoban_py_envpool_test.py"],
+    main="sokoban_py_envpool_test.py",
+    deps=[
         ":registration",
         ":sokoban",
         "//envpool",
         requirement("numpy"),
-        requirement("absl-py"),
+        requirement("pytest"),
     ],
 )
 
 pybind_extension(
-    name = "sokoban_envpool",
-    srcs = [
+    name="sokoban_envpool",
+    srcs=[
         "level_loader.cc",
         "sokoban_envpool.cc",
     ],
-    linkopts = [
+    linkopts=[
         "-ldl",
     ],
-    deps = [
+    deps=[
         ":sokoban_envpool_h",
         "//envpool/core:py_envpool",
     ],

envpool/sokoban/level_loader.cc

@@ -23,20 +23,28 @@
 #include <stdexcept>
 #include <string>
 
+#include "envpool/sokoban/utils.h"
+
 namespace sokoban {
 
-LevelLoader::LevelLoader(const std::filesystem::path& base_path, int verbose)
-    : levels_(0),
+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();
 }
 
 static const std::array<char, kMaxLevelObject + 1> kPrintLevelKey{
-    '#', ' ', '.', 'a', '@', '$', 's'};
+    '#', ' ', '.', 'a', '$', '@', 's'};
 
 void AddLine(SokobanLevel& level, const std::string& line) {
   auto start = line.at(0);
@@ -89,11 +97,19 @@ void PrintLevel(std::ostream& os, const 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 {
+    const size_t load_file_idx = SafeUniformInt(
+        static_cast<size_t>(0), level_file_paths_.size() - 1, gen);
+    file_path = level_file_paths_.at(load_file_idx);
+  }
   std::ifstream file(file_path);
 
   levels_.clear();
@@ -134,15 +150,17 @@ 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
+    std::cout << "***Loaded " << levels_.size() << " levels from " << file_path
               << std::endl;
     if (verbose >= 2) {
       PrintLevel(std::cout, levels_.at(0));
@@ -153,17 +171,20 @@ void LevelLoader::LoadNewFile(std::mt19937& gen) {
   }
 }
 
-std::vector<SokobanLevel>::iterator LevelLoader::RandomLevel(
-    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.");
+  }
   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

@@ -36,16 +36,22 @@ constexpr uint8_t kMaxLevelObject = kPlayerOnTarget;
 
 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;
 
-  std::vector<SokobanLevel>::iterator RandomLevel(std::mt19937& gen);
-  explicit LevelLoader(const std::filesystem::path& base_path, int verbose = 0);
+  std::vector<SokobanLevel>::iterator GetLevel(std::mt19937& gen);
+  explicit LevelLoader(const std::filesystem::path& base_path,
+                       bool load_sequentially, int n_levels_to_load,
+                       int verbose = 0);
 };
 
 void PrintLevel(std::ostream& os, const SokobanLevel& vec);

envpool/sokoban/registration.py

@@ -23,4 +23,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

@@ -20,17 +20,17 @@
 #include <vector>
 
 #include "envpool/core/py_envpool.h"
+#include "envpool/sokoban/utils.h"
 
 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);
-  current_max_episode_steps_ = episode_length_rand(gen_);
+  current_max_episode_steps_ =
+      SafeUniformInt(min_episode_steps, max_episode_steps, 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()="

envpool/sokoban/sokoban_envpool.h

@@ -46,7 +46,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) {
@@ -71,7 +73,9 @@ 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_(kWall, static_cast<std::size_t>(dim_room_ * dim_room_)),
         verbose_(static_cast<int>(spec.config["verbose"_])),
         current_max_episode_steps_(