change nle language wrapper so we can set render_mode

Author: BartekCupial

balrog/config/config.yaml

@@ -41,6 +41,7 @@ envs:
   names: babyai-babaisai-textworld-crafter-nle-minihack   # Environments to evaluate, separated by hyphens
   env_kwargs:
     seed: null                # Random seed; null means a random seed is used
+    render_mode: null
   nle_kwargs:
     character: "@"            # Character representing the agent in NLE
     max_episode_steps: 100_000  # Max steps per episode in NLE

balrog/environments/minihack/minihack_env.py

@@ -25,6 +25,7 @@ def make_minihack_env(env_name, task, config, render_mode: Optional[str] = None)
             "tty_colors",
         ],
         **minihack_kwargs,
+        render_mode=render_mode,
     )
     if skip_more:
         env = AutoMore(env)

balrog/environments/nle/base.py

@@ -102,6 +102,16 @@ def step(self, action):
     def post_step(self, nle_obs):
         return self.nle_process_obs(nle_obs)
 
+    def render(self):
+        mode = self.env.render_mode
+
+        if mode == "tiles":
+            return self.render_tiles_from_obs(self.env.unwrapped._last_obs)
+        elif mode == "tty_image":
+            return self.render_tty_from_obs(self.env.unwrapped._last_obs)
+        else:
+            return self.env.render()
+
     @property
     def default_action(self):
         if "minihack" in self.env.spec.id.lower():
@@ -113,7 +123,7 @@ def get_text_action(self, action):
         return self.action_str_enum_map[action]
 
     def nle_process_obs(self, nle_obs):
-        img = Image.fromarray(self.render("tiles")).convert("RGB") if self.vlm else None
+        img = Image.fromarray(self.render_tiles_from_obs(nle_obs)).convert("RGB") if self.vlm else None
         text = self.nle_obs_type(nle_obs)
 
         return {
@@ -131,15 +141,14 @@ def nle_obs_type(self, nle_obs):
         else:
             raise ValueError(f'"{self.prompt_mode}" is not a valid prompt mode.')
 
-    def render(self, mode="human"):
-        if mode in ("tiles", "tty_image"):
-            obs = self.env.unwrapped.last_observation
-            key_idx = self.env.unwrapped._observation_keys.index
-            if mode == "tiles":
-                return rgb_render_image(obs[key_idx("glyphs")])
-            else:
-                return tty_render_image(obs[key_idx("tty_chars")], obs[key_idx("tty_colors")])
-        return self.env.render(mode)
+    def render_tiles_from_obs(self, obs):
+        # Custom tiles rendering from latest observation
+        key_idx = self.env.unwrapped._observation_keys.index
+        return rgb_render_image(obs[key_idx("glyphs")])
+
+    def render_tty_from_obs(self, obs):
+        key_idx = self.env.unwrapped._observation_keys.index
+        return tty_render_image(obs[key_idx("tty_chars")], obs[key_idx("tty_colors")])
 
     def get_stats(self):
         return self.progress.__dict__

balrog/environments/nle/nle_env.py

@@ -14,7 +14,7 @@ def make_nle_env(env_name, task, config, render_mode: Optional[str] = None):
     nle_kwargs = dict(config.envs.nle_kwargs)
     skip_more = nle_kwargs.pop("skip_more", False)
     vlm = True if config.agent.max_image_history > 0 else False
-    env = gym.make(task, **nle_kwargs)
+    env = gym.make(task, **nle_kwargs, render_mode=render_mode)
     if skip_more:
         env = AutoMore(env)
 

balrog/evaluator.py

@@ -254,7 +254,7 @@ def run_episode(self, task, agent, process_num=None, position=0, episode_idx=0):
         Returns:
             dict: Log of the episode containing statistics and results.
         """
-        env = make_env(self.env_name, task, self.config)
+        env = make_env(self.env_name, task, self.config, render_mode=self.config.envs.env_kwargs.render_mode)
         agent.reset()
 
         seed = self.config.envs.env_kwargs.seed