Skip to content
This repository was archived by the owner on May 6, 2024. It is now read-only.

Commit c32a205

Browse files
committed
Add notebook rendering style, that sets background as black.
1 parent d62db5c commit c32a205

File tree

1 file changed

+6
-4
lines changed

1 file changed

+6
-4
lines changed

nle/env/base.py

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -541,14 +541,15 @@ def get_seeds(self):
541541
"""
542542
return self.env.get_current_seeds()
543543

544-
def get_tty_rendering(self, tty_chars, tty_colors):
544+
def get_tty_rendering(self, tty_chars, tty_colors, set_background=False):
545545
# TODO: Merge with "full" rendering below. Why do we have both?
546+
ansi_code = "\033[%d;3%d;40m%s" if set_background else "\033[%d;3%dm%s"
546547
rows, cols = tty_chars.shape
547548
result = ""
548549
for i in range(rows):
549550
line = "\n"
550551
for j in range(cols):
551-
line += "\033[%d;3%dm%s" % (
552+
line += ansi_code % (
552553
bool(tty_colors[i, j] & TTY_BRIGHT),
553554
tty_colors[i, j] & ~TTY_BRIGHT,
554555
chr(tty_chars[i, j]),
@@ -559,13 +560,14 @@ def get_tty_rendering(self, tty_chars, tty_colors):
559560
def render(self, mode="human"):
560561
"""Renders the state of the environment."""
561562
chars_index = self._observation_keys.index("chars")
562-
if mode == "human":
563+
if mode == "human" or mode == "notebook":
564+
set_background = mode == "notebook"
563565
obs = self.last_observation
564566
tty_chars_index = self._observation_keys.index("tty_chars")
565567
tty_colors_index = self._observation_keys.index("tty_colors")
566568
tty_chars = obs[tty_chars_index]
567569
tty_colors = obs[tty_colors_index]
568-
rendering = self.get_tty_rendering(tty_chars, tty_colors)
570+
rendering = self.get_tty_rendering(tty_chars, tty_colors, set_background)
569571
print(rendering)
570572
return
571573
elif mode == "full":

0 commit comments

Comments
 (0)