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

Commit e55f232

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

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
@@ -515,14 +515,15 @@ def get_seeds(self):
515515
"""
516516
return self.env.get_current_seeds()
517517

518-
def get_tty_rendering(self, tty_chars, tty_colors):
518+
def get_tty_rendering(self, tty_chars, tty_colors, set_background=False):
519519
# TODO: Merge with "full" rendering below. Why do we have both?
520+
ansi_code = "\033[%d;3%d;40m%s" if set_background else "\033[%d;3%dm%s"
520521
rows, cols = tty_chars.shape
521522
result = ""
522523
for i in range(rows):
523524
line = "\n"
524525
for j in range(cols):
525-
line += "\033[%d;3%dm%s" % (
526+
line += ansi_code % (
526527
bool(tty_colors[i, j] & TTY_BRIGHT),
527528
tty_colors[i, j] & ~TTY_BRIGHT,
528529
chr(tty_chars[i, j]),
@@ -533,13 +534,14 @@ def get_tty_rendering(self, tty_chars, tty_colors):
533534
def render(self, mode="human"):
534535
"""Renders the state of the environment."""
535536
chars_index = self._observation_keys.index("chars")
536-
if mode == "human":
537+
if mode == "human" or mode == "notebook":
538+
set_background = mode == "notebook"
537539
obs = self.last_observation
538540
tty_chars_index = self._observation_keys.index("tty_chars")
539541
tty_colors_index = self._observation_keys.index("tty_colors")
540542
tty_chars = obs[tty_chars_index]
541543
tty_colors = obs[tty_colors_index]
542-
rendering = self.get_tty_rendering(tty_chars, tty_colors)
544+
rendering = self.get_tty_rendering(tty_chars, tty_colors, set_background)
543545
print(rendering)
544546
return
545547
elif mode == "full":

0 commit comments

Comments
 (0)