change saving and loading

Author: BartekCupial

nle/env/base.py

@@ -193,8 +193,6 @@ def __init__(
         allow_all_yn_questions=False,
         allow_all_modes=False,
         spawn_monsters=True,
-        gamesavedir=None,
-        gameloaddir=None,
     ):
         """Constructs a new NLE environment.
 
@@ -310,8 +308,6 @@ def __init__(
             wizard=wizard,
             spawn_monsters=spawn_monsters,
             scoreprefix=scoreprefix,
-            gamesavedir=gamesavedir,
-            gameloaddir=gameloaddir,
         )
         self._close_nethack = weakref.finalize(self, self.nethack.close)
 
@@ -521,8 +517,25 @@ def render(self, mode="human"):
 
         return super().render(mode=mode)
 
-    def save(self, gamesavedir=None):
-        return self.nethack.save(gamesavedir=gamesavedir)
+    def save(self, gamesavedir):
+        # save calls function `dosave0` (responsible for saving the game) and copies save to `gamesavedir`
+        # we have to call this function explicitly because action combination ord(S) + ord(y) are overriden 
+        self.nethack.save(gamesavedir)
+
+        # we use reset to turn on the game, since saving turns it off
+        # TODO: can we dump save files without turning off the game??? 
+        self.last_observation = self.nethack.reset(self.ttyrec)
+        
+        # get out of the menu after we load new game
+        # TODO: we have one more problem, after resetting the game we have left all open menus (like inventory) 
+        # additionally we see new message "Welcome back to NetHack"
+        self.nethack.step(ASCII_ESC)       
+
+    def load(self, gameloaddir):
+        self.last_observation = self.nethack.load(gameloaddir)
+
+        # get out of the menu after we load new game
+        self.nethack.step(ASCII_ESC)
 
     def __repr__(self):
         return "<%s>" % self.__class__.__name__

nle/nethack/nethack.py

@@ -166,11 +166,7 @@ def __init__(
         hackdir=HACKDIR,
         spawn_monsters=True,
         scoreprefix="",
-        gamesavedir=None,
-        gameloaddir=None,
     ):
-        self.gamesavedir = gamesavedir
-        self.gameloaddir = gameloaddir
         self._copy = copy
 
         if not os.path.exists(hackdir) or not os.path.exists(
@@ -184,31 +180,24 @@ def __init__(
         self._tempdir = tempfile.TemporaryDirectory(prefix="nle")
         self._vardir = self._tempdir.name
 
-        if self.gameloaddir:
-            # restore files (save) from directory
-            shutil.copytree(self.gameloaddir, self._vardir, dirs_exist_ok=True)
+        # Symlink a nhdat.
+        os.symlink(os.path.join(hackdir, "nhdat"), os.path.join(self._vardir, "nhdat"))
 
-            self.dlpath = os.path.join(self._vardir, "libnethack.so")
-            self._dl = open(self.dlpath, "r")
+        # Touch files, so lock_file() in files.c passes.
+        for fn in ["perm", "record", "logfile"]:
+            os.close(os.open(os.path.join(self._vardir, fn), os.O_CREAT))
+        if scoreprefix:
+            os.close(os.open(scoreprefix + "xlogfile", os.O_CREAT))
         else:
-            # Symlink a nhdat.
-            os.symlink(os.path.join(hackdir, "nhdat"), os.path.join(self._vardir, "nhdat"))
-
-            # Touch files, so lock_file() in files.c passes.
-            for fn in ["perm", "record", "logfile"]:
-                os.close(os.open(os.path.join(self._vardir, fn), os.O_CREAT))
-            if scoreprefix:
-                os.close(os.open(scoreprefix + "xlogfile", os.O_CREAT))
-            else:
-                os.close(os.open(os.path.join(self._vardir, "xlogfile"), os.O_CREAT))
+            os.close(os.open(os.path.join(self._vardir, "xlogfile"), os.O_CREAT))
 
-            os.mkdir(os.path.join(self._vardir, "save"))
+        os.mkdir(os.path.join(self._vardir, "save"))
 
-            # An assortment of hacks:
-            # Copy our .so into self._vardir to load several copies of the dl.
-            # (Or use a memfd_create hack to create a file that gets deleted on
-            # process exit.)
-            self._dl, self.dlpath = _new_dl(self._vardir)
+        # An assortment of hacks:
+        # Copy our .so into self._vardir to load several copies of the dl.
+        # (Or use a memfd_create hack to create a file that gets deleted on
+        # process exit.)
+        self._dl, self.dlpath = _new_dl(self._vardir)
 
         # Finalize even when the rest of this constructor fails.
         self._finalizer = weakref.finalize(self, _close, None, self._dl, self._tempdir)
@@ -332,14 +321,19 @@ def in_normal_game(self):
     def how_done(self):
         return self._pynethack.how_done()
 
-    def save(self, gamesavedir=None):
-        if gamesavedir:
-            savedir = gamesavedir
-        else:
-            savedir = self.gamesavedir
-
-        assert savedir is not None
+    def save(self, gamesavedir):
+        assert gamesavedir is not None
+       
+        self._pynethack.save()
+        shutil.copytree(self._vardir, gamesavedir, dirs_exist_ok=True)
+        self._pynethack.set_use_seed_init(True)
+
+    def load(self, gameloaddir):
+        assert gameloaddir is not None
+        
+        self._pynethack.end()
+        shutil.copytree(gameloaddir, self._vardir, dirs_exist_ok=True)
+        self._pynethack.set_use_seed_init(True)
+        self._pynethack.start()
 
-        success = self._pynethack.save()
-        shutil.copytree(self._vardir, savedir, dirs_exist_ok=True)
-        return success
+        return self._step_return()

sys/unix/nledl.c

@@ -156,8 +156,7 @@ int
 nle_save(nledl_ctx *nledl)
 {   
     int success;
-    void *(*dosave0)();
-
+    int *(*dosave0)();
     dosave0 = dlsym(nledl->dlhandle, "dosave0");
     success = dosave0();
 

win/rl/pynethack.cc

@@ -101,10 +101,6 @@ checked_conversion(py::handle h, const std::vector<ssize_t> &shape)
 class Nethack
 {
   public:
-    int save() {
-        return nle_save(nle_);
-    }
-
     Nethack(std::string dlpath, std::string ttyrec, std::string hackdir,
             std::string nethackoptions, bool spawn_monsters,
             std::string scoreprefix)
@@ -205,6 +201,42 @@ class Nethack
         ttyrec_ = f;
     }
 
+    int save() 
+    {
+        int success;
+        success = nle_save(nle_);        
+        return success;
+    }
+
+    void
+    set_use_seed_init(bool use_seed_init) 
+    {
+        use_seed_init_ = use_seed_init;
+    }
+
+    void 
+    end()
+    {
+        // We use this function for loading the game from save
+        // 1. end()
+        // 2. we copy save to hackdir in `Nethack`
+        // 3. start()
+        if (nle_) {
+            nle_end(nle_);
+            nle_ = nullptr;
+        }
+    }
+
+    void
+    start() 
+    {
+        py::gil_scoped_release gil;
+
+        nle_ = 
+            nle_start(dlpath_.c_str(), &obs_, ttyrec_,
+                          use_seed_init_ ? &seed_init_ : nullptr, &settings_);
+    }
+
     void
     set_buffers(py::object glyphs, py::object chars, py::object colors,
                 py::object specials, py::object blstats, py::object message,
@@ -282,7 +314,7 @@ class Nethack
         seed_init_.seeds[0] = core;
         seed_init_.seeds[1] = disp;
         seed_init_.reseed = reseed;
-        use_seed_init = true;
+        use_seed_init_ = true;
 #else
         throw std::runtime_error("Seeding not enabled");
 #endif
@@ -351,11 +383,11 @@ class Nethack
         if (!nle_) {
             nle_ =
                 nle_start(dlpath_.c_str(), &obs_, ttyrec ? ttyrec : ttyrec_,
-                          use_seed_init ? &seed_init_ : nullptr, &settings_);
+                          use_seed_init_ ? &seed_init_ : nullptr, &settings_);
         } else
             nle_reset(nle_, &obs_, ttyrec,
-                      use_seed_init ? &seed_init_ : nullptr, &settings_);
-        use_seed_init = false;
+                      use_seed_init_ ? &seed_init_ : nullptr, &settings_);
+        use_seed_init_ = false;
 
         if (obs_.done)
             throw std::runtime_error("NetHack done right after reset");
@@ -365,7 +397,7 @@ class Nethack
     nle_obs obs_;
     std::vector<py::object> py_buffers_;
     nle_seeds_init_t seed_init_;
-    bool use_seed_init = false;
+    bool use_seed_init_ = false;
     nledl_ctx *nle_ = nullptr;
     std::FILE *ttyrec_ = nullptr;
     nle_settings settings_;
@@ -409,6 +441,9 @@ PYBIND11_MODULE(_pynethack, m)
         .def("in_normal_game", &Nethack::in_normal_game)
         .def("how_done", &Nethack::how_done)
         .def("save", &Nethack::save)
+        .def("start", &Nethack::start)
+        .def("end", &Nethack::end)
+        .def("set_use_seed_init", &Nethack::set_use_seed_init, py::arg("use_seed_init"))
         .def("set_wizkit", &Nethack::set_wizkit);
 
     py::module mn = m.def_submodule(