Fix terminal size at 80x24, write bzip2'ed ttyrecs

.circleci/config.yml

@@ -54,7 +54,7 @@ commands:
           command: |
             sudo apt-get install software-properties-common
             sudo apt-get update -y
-            sudo apt-get install -y libncurses5 libncurses5-dev flex bison cmake
+            sudo apt-get install -y libncurses5 libncurses5-dev flex bison cmake libbz2-dev
 
   pip-install-nle:
     description: Installs NLE and its dependencies

CMakeLists.txt

@@ -93,7 +93,7 @@ target_link_directories(nethack PUBLIC /usr/local/lib)
 
 # On Ubuntu 20.04, libncurses6 seems to cause a SEGFAULT on tgetent?
 find_library(NCURSES_LIB NAMES libncurses.so.5 libncurses ncurses)
-target_link_libraries(nethack PUBLIC m fcontext ${NCURSES_LIB})
+target_link_libraries(nethack PUBLIC m fcontext ${NCURSES_LIB} bz2)
 
 # dlopen wrapper library
 add_library(nethackdl STATIC "sys/unix/nledl.c")

README.md

@@ -44,7 +44,7 @@ can be installed by doing:
 ```bash
 # Python and most build deps
 $ sudo apt-get install -y build-essential autoconf libtool pkg-config \
-    python3-dev python3-pip python3-numpy git libncurses5-dev flex bison
+    python3-dev python3-pip python3-numpy git libncurses5-dev flex bison libbz2-dev
 
 # recent cmake version
 $ wget -O - https://apt.kitware.com/keys/kitware-archive-latest.asc 2>/dev/null | sudo apt-key add -

include/nle.h

@@ -1,6 +1,8 @@
 #ifndef NLE_H
 #define NLE_H
 
+#define NLE_BZ2_TTYRECS
+
 #include <stdio.h>
 
 #include <fcontext/fcontext.h>
@@ -21,6 +23,10 @@ typedef struct nle_globals {
     char *outbuf_write_ptr;
     char *outbuf_write_end;
 
+#ifdef NLE_BZ2_TTYRECS
+    void *ttyrec_bz2;
+#endif
+
     boolean done;
     nle_obs *observation;
 } nle_ctx_t;

include/unixconf.h

@@ -139,7 +139,7 @@
 #define TIMED_DELAY
 #endif
 
-/* #define AVOID_WIN_IOCTL */ /* ensure USE_WIN_IOCTL remains undefined */
+#define AVOID_WIN_IOCTL /* ensure USE_WIN_IOCTL remains undefined */
 
 /*
  * If you define MAIL, then the player will be notified of new mail

nle/env/base.py

@@ -201,7 +201,7 @@ def __init__(
 
         if self.savedir:
             self._ttyrec_pattern = os.path.join(
-                self.savedir, "nle.%i.%%i.ttyrec" % os.getpid()
+                self.savedir, "nle.%i.%%i.ttyrec.bz2" % os.getpid()
             )
             ttyrec = self._ttyrec_pattern % 0
         else:

nle/nethack/nethack.py

@@ -80,7 +80,7 @@ def __init__(
         self,
         observation_keys=OBSERVATION_DESC.keys(),
         playername="Agent-mon-hum-neu-mal",
-        ttyrec="nle.ttyrec",
+        ttyrec="nle.ttyrec.bz2",
         options=None,
         copy=False,
         wizard=False,

nle/scripts/play.py

@@ -74,7 +74,7 @@ def play(env, mode, ngames, max_steps, seeds, savedir, no_render, debug):
     if is_raw_env:
         if savedir is not None:
             os.makedirs(savedir, exist_ok=True)
-            ttyrec = os.path.join(savedir, "nle.ttyrec")
+            ttyrec = os.path.join(savedir, "nle.ttyrec.bz2")
         else:
             ttyrec = "/dev/null"
         env = nethack.Nethack(ttyrec=ttyrec)

nle/tests/test_envs.py

@@ -224,7 +224,7 @@ def test_rollout(self, env_name, rollout_len):
             env.close()
 
             assert os.path.exists(
-                os.path.join(savedir, "nle.%i.0.ttyrec" % os.getpid())
+                os.path.join(savedir, "nle.%i.0.ttyrec.bz2" % os.getpid())
             )
 
     def test_rollout_no_archive(self, env_name, rollout_len):

src/nle.c

@@ -1,8 +1,7 @@
 
 #include <assert.h>
-#include <sys/time.h>
-
 #include <string.h>
+#include <sys/time.h>
 
 #define NEED_VARARGS
 #include "hack.h"
@@ -11,6 +10,10 @@
 
 #include "nle.h"
 
+#ifdef NLE_BZ2_TTYRECS
+#include <bzlib.h>
+#endif
+
 #define STACK_SIZE (1 << 15) // 32KiB
 
 #ifndef __has_feature
@@ -31,6 +34,12 @@ init_nle(FILE *ttyrec)
     assert(ttyrec != NULL);
     nle->ttyrec = ttyrec;
 
+#ifdef NLE_BZ2_TTYRECS
+    int bzerror;
+    nle->ttyrec_bz2 = BZ2_bzWriteOpen(&bzerror, ttyrec, 9, 0, 0);
+    assert(bzerror == BZ_OK);
+#endif
+
     nle->outbuf_write_ptr = nle->outbuf;
     nle->outbuf_write_end = nle->outbuf + sizeof(nle->outbuf);
 
@@ -59,6 +68,20 @@ mainloop(fcontext_transfer_t ctx_transfer)
     unixmain(1, argv);
 }
 
+boolean
+write_data(void *buf, int length)
+{
+    nle_ctx_t *nle = current_nle_ctx;
+#ifdef NLE_BZ2_TTYRECS
+    int bzerror;
+    BZ2_bzWrite(&bzerror, nle->ttyrec_bz2, buf, length);
+    assert(bzerror == BZ_OK);
+#else
+    assert(fwrite(buf, 1, length, nle->ttyrec) == 0);
+#endif
+    return TRUE;
+}
+
 boolean
 write_header(int length, unsigned char channel)
 {
@@ -70,18 +93,9 @@ write_header(int length, unsigned char channel)
     buffer[1] = tv.tv_usec;
     buffer[2] = length;
 
-    nle_ctx_t *nle = current_nle_ctx;
-
     /* Assumes little endianness */
-    if (fwrite(buffer, sizeof(int), 3, nle->ttyrec) == 0) {
-        assert(FALSE);
-        return FALSE;
-    }
-
-    if (fputc((int) channel, nle->ttyrec) != (int) channel) {
-        assert(FALSE);
-        return FALSE;
-    }
+    write_data(buffer, 3 * sizeof(int));
+    write_data(&channel, 1);
 
     return TRUE;
 }
@@ -102,12 +116,17 @@ nle_fflush(FILE *stream)
     ssize_t length = nle->outbuf_write_ptr - nle->outbuf;
     if (length == 0)
         return 0;
-    /* TODO(heiner): Given that we do our own buffering, consider
-     * using file descriptors instead of the ttyrec FILE*. */
+
     write_header(length, 0);
-    fwrite(nle->outbuf, 1, length, nle->ttyrec);
+    write_data(nle->outbuf, length);
+
     nle->outbuf_write_ptr = nle->outbuf;
+
+#ifdef NLE_BZ2_TTYRECS
+    return 0;
+#else
     return fflush(nle->ttyrec);
+#endif
 }
 
 /*
@@ -219,6 +238,10 @@ init_random(int FDECL((*fn), (int) ))
 nle_ctx_t *
 nle_start(nle_obs *obs, FILE *ttyrec, nle_seeds_init_t *seed_init)
 {
+    /* Set CO and LI to control ttyrec output size. */
+    CO = 80;
+    LI = 24;
+
     nle_ctx_t *nle = init_nle(ttyrec);
     nle->observation = obs;
     nle_seeds_init = seed_init;
@@ -244,7 +267,7 @@ nle_step(nle_ctx_t *nle, nle_obs *obs)
     current_nle_ctx = nle;
     nle->observation = obs;
     write_header(1, 1);
-    fputc(obs->action, nle->ttyrec);
+    write_data(&obs->action, 1);
     fcontext_transfer_t t = jump_fcontext(nle->generatorcontext, obs);
     nle->generatorcontext = t.ctx;
     nle->done = (t.data == NULL);
@@ -256,7 +279,6 @@ nle_step(nle_ctx_t *nle, nle_obs *obs)
 void
 nle_end(nle_ctx_t *nle)
 {
-    nle_fflush(stdout);
     if (!nle->done) {
         /* Reset without closing nethack. Need free memory, etc.
          * this is what nh_terminate in end.c does. I hope it's enough. */
@@ -265,6 +287,13 @@ nle_end(nle_ctx_t *nle)
             dlb_cleanup();
         }
     }
+    nle_fflush(stdout);
+
+#ifdef NLE_BZ2_TTYRECS
+    int bzerror;
+    BZ2_bzWriteClose(&bzerror, nle->ttyrec_bz2, 0, NULL, NULL);
+    assert(bzerror == BZ_OK);
+#endif
 
     destroy_fcontext_stack(&nle->stack);
     free(nle);