image obs rendering

Author: Howuhh

src/xminigrid/__init__.py

@@ -2,7 +2,7 @@
 from .registration import make, register, registered_environments
 
 # TODO: add __all__
-__version__ = "0.6.0"
+__version__ = "0.7.0"
 
 # ---------- XLand-MiniGrid environments ----------
 

src/xminigrid/core/constants.py

@@ -1,48 +1,44 @@
 import jax.numpy as jnp
 from flax import struct
 
+NUM_ACTIONS = 6
+
 # GRID: [tile, color]
 NUM_LAYERS = 2
-NUM_TILES = 15
-NUM_COLORS = 14
-NUM_ACTIONS = 6
+NUM_TILES = 13
+NUM_COLORS = 12
 
 
-# TODO: do we really need END_OF_MAP? seem like unseen can be used instead...
 # enums, kinda...
 class Tiles(struct.PyTreeNode):
     EMPTY: int = struct.field(pytree_node=False, default=0)
-    END_OF_MAP: int = struct.field(pytree_node=False, default=1)
-    UNSEEN: int = struct.field(pytree_node=False, default=2)
-    FLOOR: int = struct.field(pytree_node=False, default=3)
-    WALL: int = struct.field(pytree_node=False, default=4)
-    BALL: int = struct.field(pytree_node=False, default=5)
-    SQUARE: int = struct.field(pytree_node=False, default=6)
-    PYRAMID: int = struct.field(pytree_node=False, default=7)
-    GOAL: int = struct.field(pytree_node=False, default=8)
-    KEY: int = struct.field(pytree_node=False, default=9)
-    DOOR_LOCKED: int = struct.field(pytree_node=False, default=10)
-    DOOR_CLOSED: int = struct.field(pytree_node=False, default=11)
-    DOOR_OPEN: int = struct.field(pytree_node=False, default=12)
-    HEX: int = struct.field(pytree_node=False, default=13)
-    STAR: int = struct.field(pytree_node=False, default=14)
+    FLOOR: int = struct.field(pytree_node=False, default=1)
+    WALL: int = struct.field(pytree_node=False, default=2)
+    BALL: int = struct.field(pytree_node=False, default=3)
+    SQUARE: int = struct.field(pytree_node=False, default=4)
+    PYRAMID: int = struct.field(pytree_node=False, default=5)
+    GOAL: int = struct.field(pytree_node=False, default=6)
+    KEY: int = struct.field(pytree_node=False, default=7)
+    DOOR_LOCKED: int = struct.field(pytree_node=False, default=8)
+    DOOR_CLOSED: int = struct.field(pytree_node=False, default=9)
+    DOOR_OPEN: int = struct.field(pytree_node=False, default=10)
+    HEX: int = struct.field(pytree_node=False, default=11)
+    STAR: int = struct.field(pytree_node=False, default=12)
 
 
 class Colors(struct.PyTreeNode):
     EMPTY: int = struct.field(pytree_node=False, default=0)
-    END_OF_MAP: int = struct.field(pytree_node=False, default=1)
-    UNSEEN: int = struct.field(pytree_node=False, default=2)
-    RED: int = struct.field(pytree_node=False, default=3)
-    GREEN: int = struct.field(pytree_node=False, default=4)
-    BLUE: int = struct.field(pytree_node=False, default=5)
-    PURPLE: int = struct.field(pytree_node=False, default=6)
-    YELLOW: int = struct.field(pytree_node=False, default=7)
-    GREY: int = struct.field(pytree_node=False, default=8)
-    BLACK: int = struct.field(pytree_node=False, default=9)
-    ORANGE: int = struct.field(pytree_node=False, default=10)
-    WHITE: int = struct.field(pytree_node=False, default=11)
-    BROWN: int = struct.field(pytree_node=False, default=12)
-    PINK: int = struct.field(pytree_node=False, default=13)
+    RED: int = struct.field(pytree_node=False, default=1)
+    GREEN: int = struct.field(pytree_node=False, default=2)
+    BLUE: int = struct.field(pytree_node=False, default=3)
+    PURPLE: int = struct.field(pytree_node=False, default=4)
+    YELLOW: int = struct.field(pytree_node=False, default=5)
+    GREY: int = struct.field(pytree_node=False, default=6)
+    BLACK: int = struct.field(pytree_node=False, default=7)
+    ORANGE: int = struct.field(pytree_node=False, default=8)
+    WHITE: int = struct.field(pytree_node=False, default=9)
+    BROWN: int = struct.field(pytree_node=False, default=10)
+    PINK: int = struct.field(pytree_node=False, default=11)
 
 
 # Only ~100 combinations so far, better to preallocate them
@@ -65,7 +61,6 @@ class Colors(struct.PyTreeNode):
 
 WALKABLE = jnp.array(
     (
-        Tiles.EMPTY,
         Tiles.FLOOR,
         Tiles.GOAL,
         Tiles.DOOR_OPEN,
@@ -83,12 +78,7 @@ class Colors(struct.PyTreeNode):
     )
 )
 
-FREE_TO_PUT_DOWN = jnp.array(
-    (
-        Tiles.EMPTY,
-        Tiles.FLOOR,
-    )
-)
+FREE_TO_PUT_DOWN = jnp.array((Tiles.FLOOR,))
 
 LOS_BLOCKING = jnp.array(
     (

src/xminigrid/core/grid.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Callable, Union
+from typing import Callable
 
 import jax
 import jax.numpy as jnp
@@ -22,7 +22,7 @@ def equal(tile1: Tile, tile2: Tile) -> Tile:
 
 def get_neighbouring_tiles(grid: GridState, y: IntOrArray, x: IntOrArray) -> tuple[Tile, Tile, Tile, Tile]:
     # end_of_map = TILES_REGISTRY[Tiles.END_OF_MAP, Colors.END_OF_MAP]
-    end_of_map = Tiles.END_OF_MAP
+    end_of_map = Tiles.EMPTY
 
     up_tile = grid.at[y - 1, x].get(mode="fill", fill_value=end_of_map)
     right_tile = grid.at[y, x + 1].get(mode="fill", fill_value=end_of_map)

src/xminigrid/core/observation.py

@@ -12,7 +12,7 @@ def crop_field_of_view(grid: GridState, agent: AgentState, height: int, width: i
     grid = jnp.pad(
         grid,
         pad_width=((height, height), (width, width), (0, 0)),
-        constant_values=Tiles.END_OF_MAP,
+        constant_values=Tiles.EMPTY,
     )
     # account for padding
     y = agent.position[0] + height
@@ -110,8 +110,8 @@ def minigrid_field_of_view(grid: GridState, agent: AgentState, height: int, widt
     fov_grid = crop_field_of_view(grid, agent, height, width)
     fov_grid = align_with_up(fov_grid, agent.direction)
     mask = generate_viz_mask_minigrid(fov_grid)
-    # set UNSEEN value for all layers (including colors, as UNSEEN color has same id value)
-    fov_grid = jnp.where(mask[..., None], fov_grid, Tiles.UNSEEN)
+    # set EMPTY as unseen value for all layers (including colors, as EMPTY color has same id value)
+    fov_grid = jnp.where(mask[..., None], fov_grid, Tiles.EMPTY)
 
     # TODO: should we even do this? Agent with good memory can remember what he picked up.
     # WARN: this can overwrite tile the agent is on, GOAL for example.

src/xminigrid/envs/xland.py

@@ -27,7 +27,6 @@
     init_tiles=jnp.array(((TILES_REGISTRY[Tiles.EMPTY, Colors.EMPTY],))),
 )
 
-_empty_tile = TILES_REGISTRY[Tiles.EMPTY, Colors.EMPTY]
 _wall_tile = TILES_REGISTRY[Tiles.WALL, Colors.GREY]
 # colors for doors between rooms
 _allowed_colors = jnp.array(

src/xminigrid/experimental/img_obs.py

@@ -0,0 +1,90 @@
+# jit-compatible RGB observations. Currently experimental!
+# if it proves useful and necessary in the future, I will consider rewriting env.render in such style also
+from __future__ import annotations
+
+import os
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from ..benchmarks import load_bz2_pickle, save_bz2_pickle
+from ..core.constants import NUM_COLORS, NUM_LAYERS, TILES_REGISTRY
+from ..rendering.rgb_render import render_tile
+from ..wrappers import Wrapper
+
+CACHE_PATH = os.environ.get("XLAND_MINIGRID_CACHE", os.path.expanduser("~/.xland_minigrid"))
+
+
+def build_cache(tiles: np.ndarray, tile_size: float = 32) -> tuple[np.ndarray, np.ndarray]:
+    cache = np.full((tiles.shape[0], tiles.shape[1], tile_size, tile_size, 3), dtype=np.uint8, fill_value=-1)
+    agent_cache = np.full((tiles.shape[0], tiles.shape[1], tile_size, tile_size, 3), dtype=np.uint8, fill_value=-1)
+
+    for y in range(tiles.shape[0]):
+        for x in range(tiles.shape[1]):
+            # rendering tile
+            tile_img = render_tile(
+                tile=tuple(tiles[y, x]),
+                agent_direction=None,
+                highlight=False,
+                tile_size=int(tile_size),
+            )
+            cache[y, x] = tile_img
+
+            # rendering agent on top
+            tile_w_agent_img = render_tile(
+                tile=tuple(tiles[y, x]),
+                agent_direction=0,
+                highlight=False,
+                tile_size=int(tile_size),
+            )
+            agent_cache[y, x] = tile_w_agent_img
+
+    return cache, agent_cache
+
+
+# building cache of pre-rendered tiles
+TILE_SIZE = 32
+
+cache_path = os.path.join(CACHE_PATH, "render_cache")
+if not os.path.exists(cache_path):
+    TILE_CACHE, TILE_W_AGENT_CACHE = build_cache(np.asarray(TILES_REGISTRY), tile_size=TILE_SIZE)
+    TILE_CACHE = jnp.asarray(TILE_CACHE).reshape(-1, TILE_SIZE, TILE_SIZE, 3)
+    TILE_W_AGENT_CACHE = jnp.asarray(TILE_W_AGENT_CACHE).reshape(-1, TILE_SIZE, TILE_SIZE, 3)
+
+    save_bz2_pickle({"tile_cache": TILE_CACHE, "tile_agent_cache": TILE_W_AGENT_CACHE}, cache_path)
+
+TILE_CACHE = load_bz2_pickle(cache_path)["tile_cache"]
+TILE_W_AGENT_CACHE = load_bz2_pickle(cache_path)["tile_agent_cache"]
+
+
+# rendering with cached tiles
+def _render_obs(obs: jax.Array) -> jax.Array:
+    view_size = obs.shape[0]
+
+    obs_flat_idxs = obs[:, :, 0] * NUM_COLORS + obs[:, :, 1]
+    # render all tiles
+    rendered_obs = jnp.take(TILE_CACHE, obs_flat_idxs, axis=0)
+
+    # add agent tile
+    agent_tile = TILE_W_AGENT_CACHE[obs_flat_idxs[view_size - 1, view_size // 2]]
+    rendered_obs = rendered_obs.at[view_size - 1, view_size // 2].set(agent_tile)
+    # [view_size, view_size, tile_size, tile_size, 3] -> [view_size * tile_size, view_size * tile_size, 3]
+    rendered_obs = rendered_obs.transpose((0, 2, 1, 3, 4)).reshape(view_size * TILE_SIZE, view_size * TILE_SIZE, 3)
+
+    return rendered_obs
+
+
+class RGBImgObservationWrapper(Wrapper):
+    def observation_shape(self, params):
+        return params.view_size * TILE_SIZE, params.view_size * TILE_SIZE, NUM_LAYERS
+
+    def reset(self, params, key):
+        timestep = self._env.reset(params, key)
+        timestep = timestep.replace(observation=_render_obs(timestep.observation))
+        return timestep
+
+    def step(self, params, timestep, action):
+        timestep = self._env.step(params, timestep, action)
+        timestep = timestep.replace(observation=_render_obs(timestep.observation))
+        return timestep

src/xminigrid/manual_control.py

@@ -14,9 +14,10 @@
 
 
 class ManualControl:
-    def __init__(self, env: Environment[EnvParamsT, EnvCarryT], env_params: EnvParamsT):
+    def __init__(self, env: Environment[EnvParamsT, EnvCarryT], env_params: EnvParamsT, agent_view: bool = False):
         self.env = env
         self.env_params = env_params
+        self.agent_view = agent_view
 
         self._reset = jax.jit(self.env.reset)
         self._step = jax.jit(self.env.step)
@@ -33,7 +34,10 @@ def __init__(self, env: Environment[EnvParamsT, EnvCarryT], env_params: EnvParam
     def render(self) -> None:
         assert self.timestep is not None
 
-        img = self.env.render(self.env_params, self.timestep)
+        if self.agent_view:
+            img = self.timestep.observation
+        else:
+            img = self.env.render(self.env_params, self.timestep)
         # [h, w, c] -> [w, h, c]
         img = np.transpose(img, axes=(1, 0, 2))
 
@@ -144,14 +148,20 @@ def close(self) -> None:
     parser.add_argument("--env-id", type=str, default="MiniGrid-Empty-5x5", choices=xminigrid.registered_environments())
     parser.add_argument("--benchmark-id", type=str, default="trivial-1m", choices=xminigrid.registered_benchmarks())
     parser.add_argument("--ruleset-id", type=int, default=0)
+    parser.add_argument("--agent-view", action="store_true")
 
     args = parser.parse_args()
     env, env_params = xminigrid.make(args.env_id)
     env = GymAutoResetWrapper(env)
 
+    if args.agent_view:
+        from xminigrid.experimental.img_obs import RGBImgObservationWrapper
+
+        env = RGBImgObservationWrapper(env)
+
     if "XLand" in args.env_id:
         bench = xminigrid.load_benchmark(args.benchmark_id)
         env_params = env_params.replace(ruleset=bench.get_ruleset(args.ruleset_id))
 
-    control = ManualControl(env=env, env_params=env_params)
+    control = ManualControl(env=env, env_params=env_params, agent_view=args.agent_view)
     control.start()

src/xminigrid/rendering/rgb_render.py

@@ -5,7 +5,7 @@
 import numpy as np
 
 from ..core.constants import Colors, Tiles
-from ..types import AgentState, GridState, IntOrArray
+from ..types import AgentState, IntOrArray
 from .utils import (
     downsample,
     fill_coords,
@@ -18,6 +18,7 @@
 )
 
 COLORS_MAP = {
+    Colors.EMPTY: np.array((255, 255, 255)),  # just a placeholder
     Colors.RED: np.array((255, 0, 0)),
     Colors.GREEN: np.array((0, 255, 0)),
     Colors.BLUE: np.array((0, 0, 255)),
@@ -32,6 +33,16 @@
 }
 
 
+def _render_empty(img: np.ndarray, color: int):
+    fill_coords(img, point_in_rect(0.45, 0.55, 0.2, 0.65), COLORS_MAP[Colors.RED])
+    fill_coords(img, point_in_rect(0.45, 0.55, 0.7, 0.85), COLORS_MAP[Colors.RED])
+
+    fill_coords(img, point_in_rect(0, 0.031, 0, 1), COLORS_MAP[Colors.RED])
+    fill_coords(img, point_in_rect(0, 1, 0, 0.031), COLORS_MAP[Colors.RED])
+    fill_coords(img, point_in_rect(1 - 0.031, 1, 0, 1), COLORS_MAP[Colors.RED])
+    fill_coords(img, point_in_rect(0, 1, 1 - 0.031, 1), COLORS_MAP[Colors.RED])
+
+
 def _render_floor(img: np.ndarray, color: int):
     # draw the grid lines (top and left edges)
     fill_coords(img, point_in_rect(0, 0.031, 0, 1), (100, 100, 100))
@@ -165,7 +176,7 @@ def _render_player(img: np.ndarray, direction: int):
     Tiles.DOOR_LOCKED: _render_door_locked,
     Tiles.DOOR_CLOSED: _render_door_closed,
     Tiles.DOOR_OPEN: _render_door_open,
-    Tiles.EMPTY: lambda img, color: img,
+    Tiles.EMPTY: _render_empty,
 }
 
 
@@ -196,7 +207,7 @@ def get_highlight_mask(grid: np.ndarray, agent: AgentState | None, view_size: in
 
 @functools.cache
 def render_tile(
-    tile: np.ndarray, agent_direction: int | None = None, highlight: bool = False, tile_size: int = 32, subdivs: int = 3
+    tile: tuple, agent_direction: int | None = None, highlight: bool = False, tile_size: int = 32, subdivs: int = 3
 ) -> np.ndarray:
     img = np.full((tile_size * subdivs, tile_size * subdivs, 3), dtype=np.uint8, fill_value=255)
     # draw tile

src/xminigrid/rendering/text_render.py

@@ -5,8 +5,6 @@
 from ..types import AgentState, RuleSet
 
 COLOR_NAMES = {
-    Colors.END_OF_MAP: "red",
-    Colors.UNSEEN: "white",
     Colors.EMPTY: "white",
     Colors.RED: "red",
     Colors.GREEN: "green",
@@ -22,8 +20,6 @@
 }
 
 TILE_STR = {
-    Tiles.END_OF_MAP: "!",
-    Tiles.UNSEEN: "?",
     Tiles.EMPTY: " ",
     Tiles.FLOOR: ".",
     Tiles.WALL: "☰",

src/xminigrid/wrappers.py

@@ -16,6 +16,12 @@ def __init__(self, env: Environment[EnvParamsT, EnvCarryT]):
     def default_params(self, **kwargs) -> EnvParamsT:
         return self._env.default_params(**kwargs)
 
+    def num_actions(self, params: EnvParamsT) -> int:
+        return self._env.num_actions(params)
+
+    def observation_shape(self, params: EnvParamsT) -> tuple[int, int, int]:
+        return self._env.observation_shape(params)
+
     def time_limit(self, params: EnvParamsT) -> int:
         return self._env.time_limit(params)