Repository: epignatelli/navix (language: python)

Issue:
Improve graphics and tiles
Maybe add themes?

We can have a `minigrid` theme, but also a different one.


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 754e9b1..64554c8 100644
--- a/README.md
+++ b/README.md
@@ -8,11 +8,16 @@
 **[Quickstart](#what-is-navix)** | **[Installation](#installation)** | **[Examples](#examples)** | **[Cite](#cite)**
 
 ## What is NAVIX?
-NAVIX is [minigrid](https://github.com/Farama-Foundation/Minigrid) in JAX, **>10000x** faster with Autograd and XLA support.
+---
+NAVIX is [minigrid](https://github.com/Farama-Foundation/Minigrid) in JAX, **>1000x** faster with Autograd and XLA support.
 You can see a superficial performance comparison [here](docs/profiling.ipynb).
 
+The library is in active development, and we are working on adding more environments and features.
+If you want join the development and contribute, please [open a discussion](https://github.com/epignatelli/navix/discussions/new?category=general) and let's have a chat!
+
 
 ## Installation
+---
 We currently support the OSs supported by JAX.
 You can find a description [here](https://github.com/google/jax#installation).
 
@@ -22,13 +27,20 @@ You might want to follow the same guide to install jax for your faviourite accel
 [TPU](https://github.com/google/jax#pip-installation-colab-tpu)
 ).
 
-Then, install `navix` and its dependencies with:
+- ### Stable
+Then, install the stable version of `navix` and its dependencies with:
 ```bash
 pip install navix
 ```
 
----
+- ### Nightly
+Or, if you prefer to install the latest version from source:
+```bash
+pip install git+https://github.com/epignatelli/navix
+```
+
 ## Examples
+---
 
 ### XLA compilation
 One straightforward use case is to accelerate the computation of the environment with XLA compilation.
@@ -64,6 +76,7 @@ TODO(epignatelli): add example.
 
 
 ## Cite
+---
 If you use `helx` please consider citing it as:
 
 ```bibtex
@@ -75,4 +88,3 @@ If you use `helx` please consider citing it as:
   journal = {GitHub repository},
   howpublished = {\url{https://github.com/epignatelli/navix}}
   }
-```
diff --git a/docs/profiling.ipynb b/docs/performance.ipynb
similarity index 61%
rename from docs/profiling.ipynb
rename to docs/performance.ipynb
index 092cabb..97618e8 100644
--- a/docs/profiling.ipynb
+++ b/docs/performance.ipynb
@@ -38,51 +38,29 @@
     "import jax.numpy as jnp\n",
     "import navix as nx\n",
     "\n",
+    "import gymnasium as gym\n",
+    "import minigrid\n",
+    "import random\n",
+    "import time\n",
     "\n",
-    "N_TIMESTEPS = 10_000\n",
+    "from timeit import timeit\n",
+    "\n",
+    "\n",
+    "N_TIMEIT_LOOPS = 5\n",
+    "N_TIMESTEPS = 100\n",
+    "N_SEEDS = 10\n",
     "\n",
     "\n",
     "def profile_navix(seed):\n",
-    "    env = nx.environments.Room(16, 16, 8)\n",
+    "    env = nx.environments.Room(16, 16, 8, observation_fn=nx.observations.rgb)\n",
     "    key = jax.random.PRNGKey(seed)\n",
     "    timestep = env.reset(key)\n",
     "    actions = jax.random.randint(key, (N_TIMESTEPS,), 0, 6)\n",
     "\n",
-    "    def body_fun(carry, x):\n",
-    "        timestep = carry\n",
-    "        action = x\n",
-    "        timestep = env.step(timestep, jnp.asarray(action))\n",
-    "        return timestep, ()\n",
-    "\n",
-    "    return jax.lax.scan(body_fun, timestep, jnp.asarray(actions, dtype=jnp.int32))[0]\n",
+    "    for i in range(N_TIMESTEPS):\n",
+    "        timestep = env.step(timestep, actions[i])\n",
     "\n",
-    "\n",
-    "f = jax.jit(jax.vmap(profile_navix))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "44SbP_tOp1dP"
-   },
-   "outputs": [],
-   "source": [
-    "# running 10_000 seeds in parallel\n",
-    "%timeit -n 5 -r 1 f(jnp.arange(10_000)).state.grid.block_until_ready()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "iSeCh6H4qRdx"
-   },
-   "outputs": [],
-   "source": [
-    "import gymnasium as gym\n",
-    "import minigrid\n",
-    "import random\n",
+    "    return timestep\n",
     "\n",
     "\n",
     "def profile_minigrid(seed):\n",
@@ -95,19 +73,31 @@
     "        if terminated or truncated:\n",
     "            observation, info = env.reset()\n",
     "    env.close()\n",
-    "    return observation"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "1b4-SrmWsYgs"
-   },
-   "outputs": [],
-   "source": [
-    "# running 1 seed\n",
-    "%timeit -n 5 -r 1 profile_minigrid(0)"
+    "    return observation\n",
+    "\n",
+    "\n",
+    "if __name__ == \"__main__\":\n",
+    "    # profile navix\n",
+    "    print(\n",
+    "        \"Profiling navix, N_SEEDS = {}, N_TIMESTEPS = {}\".format(N_SEEDS, N_TIMESTEPS)\n",
+    "    )\n",
+    "    seeds = jnp.arange(N_SEEDS)\n",
+    "\n",
+    "    print(\"\\tCompiling...\")\n",
+    "    start = time.time()\n",
+    "    f = jax.jit(jax.vmap(profile_navix)).lower(seeds).compile()\n",
+    "    print(\"\\tCompiled in {:.2f}s\".format(time.time() - start))\n",
+    "\n",
+    "    print(\"\\tRunning ...\")\n",
+    "    res_navix = timeit(\n",
+    "        lambda: f(seeds).state.grid.block_until_ready(), number=N_TIMEIT_LOOPS\n",
+    "    )\n",
+    "    print(res_navix)\n",
+    "\n",
+    "    # profile minigrid\n",
+    "    print(\"Profiling minigrid, N_SEEDS = 1, N_TIMESTEPS = {}\".format(N_TIMESTEPS))\n",
+    "    res_minigrid = timeit(lambda: profile_minigrid(0), number=N_TIMEIT_LOOPS)\n",
+    "    print(res_minigrid)"
    ]
   },
   {
@@ -132,7 +122,8 @@
    "name": "python3"
   },
   "language_info": {
-   "name": "python"
+   "name": "python",
+   "version": "3.9.16"
   }
  },
  "nbformat": 4,
diff --git a/navix/_version.py b/navix/_version.py
index 49e9dba..dc076c5 100644
--- a/navix/_version.py
+++ b/navix/_version.py
@@ -18,5 +18,5 @@
 # under the License.
 
 
-__version__ = "0.2.2"
+__version__ = "0.3.2"
 __version_info__ = tuple(int(i) for i in __version__.split(".") if i.isdigit())
diff --git a/navix/actions.py b/navix/actions.py
index 0713b23..d944ea2 100644
--- a/navix/actions.py
+++ b/navix/actions.py
@@ -24,31 +24,23 @@
 import jax.numpy as jnp
 from jax import Array
 
-from .components import Consumable, Pickable, State
+from .components import Door, Key, State, DISCARD_PILE_COORDS
+from .grid import translate, rotate
 
 
 DIRECTIONS = {0: "east", 1: "south", 2: "west", 3: "north"}
 
 
 def _rotate(state: State, spin: int) -> State:
-    direction = (state.player.direction + spin) % 4
+    direction = rotate(state.player.direction, spin)
     player = state.player.replace(direction=direction)
     return state.replace(player=player)
 
 
-def _translate(position: Array, direction: Array) -> Array:
-    moves = (
-        lambda position: position + jnp.asarray((0, 1)),  # east
-        lambda position: position + jnp.asarray((1, 0)),  # south
-        lambda position: position + jnp.asarray((0, -1)),  # west
-        lambda position: position + jnp.asarray((-1, 0)),  # north
-    )
-    return jax.lax.switch(direction, moves, position)
-
-
-def _move_allowed(state: State, position: Array) -> Array:
+def _walkable(state: State, position: Array) -> Array:
     # according to the grid
     walkable = jnp.equal(state.grid[tuple(position)], 0)
+
     # and not occupied by another non-walkable entity
     occupied_keys = jax.vmap(lambda x: jnp.array_equal(x, position))(
         state.keys.position
@@ -57,12 +49,13 @@ def _move_allowed(state: State, position: Array) -> Array:
         state.doors.position
     )
     occupied = jnp.any(jnp.concatenate([occupied_keys, occupied_doors]))
+    # return: if walkable and not occupied
     return jnp.logical_and(walkable, jnp.logical_not(occupied))
 
 
 def _move(state: State, direction: Array) -> State:
-    new_position = _translate(state.player.position, direction)
-    can_move = _move_allowed(state, new_position)
+    new_position = translate(state.player.position, direction)
+    can_move = _walkable(state, new_position)
     new_position = jnp.where(can_move, new_position, state.player.position)
     player = state.player.replace(position=new_position)
     return state.replace(player=player)
@@ -107,46 +100,52 @@ def left(state: State) -> State:
     return _move(state, state.player.direction + 3)
 
 
+def _one_many_position_equal(a: Array, b: Array) -> Array:
+    assert a.ndim == 1 and b.ndim == 2
+    is_equal = jnp.sum(a[None] - b, axis=-1) == 0
+    assert is_equal.shape == (b.shape[0],)
+    return is_equal
+
+
 def pickup(state: State) -> State:
-    position_in_front = _translate(state.player.position, state.player.direction)
-
-    def _update(key: Pickable) -> Tuple[Array, Pickable]:
-        match = jnp.array_equal(position_in_front, key.position)
-        # update player's pocket
-        pocket = jnp.where(match, key.id, state.player.pocket)
-        # set to (-1, -1) the position of the key that was picked up
-        unset_position = jnp.asarray((-1, -1))
-        position = jnp.where(match, unset_position, key.position)
-        key = key.replace(position=position)
-        return pocket, key
-
-    pockets, keys = jax.vmap(_update)(state.keys)
-    pocket = jnp.max(pockets, axis=0)
-    player = state.player.replace(pocket=pocket)
+    position_in_front = translate(state.player.position, state.player.direction)
+
+    key_found = _one_many_position_equal(position_in_front, state.keys.position)
+
+    # update keys
+    positions = jnp.where(key_found, DISCARD_PILE_COORDS, state.keys.position)
+    keys = state.keys.replace(position=positions)
+
+    # update player's pocket, if the pocket has something else, we overwrite it
+    key = jnp.sum(state.keys.id * key_found, dtype=jnp.int32)
+    player = jax.lax.cond(jnp.any(key_found), lambda: state.player.replace(pocket=key), lambda: state.player)
+
     return state.replace(player=player, keys=keys)
 
 
 def open(state: State) -> State:
-    position_in_front = _translate(state.player.position, state.player.direction)
-
-    def _update(door: Consumable) -> Tuple[Array, Consumable]:
-        match = jnp.array_equal(position_in_front, door.position)
-        replacement = jnp.asarray((match - 1) * door.replacement, dtype=jnp.int32)
-
-        # update grid
-        grid = jnp.zeros_like(state.grid).at[tuple(door.position)].set(replacement)
-
-        # set to (-1, -1) the position of the door that was opened
-        unset_position = jnp.asarray((-1, -1))
-        position = jnp.where(match, unset_position, door.position)
-        door = door.replace(position=position)
-        return grid, door
-
-    grid, doors = jax.vmap(_update)(state.doors)
-    # the max makes sure that if there was a wall (-1), and it has been opened (x>0)
-    # we get the new value of the grid
-    grid = jnp.max(grid, axis=0)
-    return state.replace(grid=grid, doors=doors)
+    # get the tile in front of the player
+    position_in_front = translate(state.player.position, state.player.direction)
+
+    # check if there is a door in front of the player
+    door_found = position_in_front[None] == state.doors.position
+    # and that, if so, either it does not require a key or the player has the key
+    requires_key = state.doors.requires != -1
+    key_match = state.player.pocket == state.doors.requires
+    can_open = door_found & (key_match | ~requires_key )
+
+    # update doors
+    # TODO(epignatelli): in the future we want to mark the door as open, instead
+    # and have a different rendering for it
+    # if the door can be opened, move it to the discard pile
+    new_positions = jnp.where(can_open, DISCARD_PILE_COORDS, state.doors.position)
+    doors = state.doors.replace(position=new_positions)
+
+    # remove key from player's pocket
+    pocket = jnp.asarray(state.player.pocket * jnp.any(can_open), dtype=jnp.int32)
+    player = jax.lax.cond(jnp.any(can_open), lambda: state.player.replace(pocket=pocket), lambda: state.player)
+
+    return state.replace(player=player, doors=doors)
 
 
 # TODO(epignatelli): a mutable dictionary here is dangerous
diff --git a/navix/components.py b/navix/components.py
index 1b306c9..3f9be09 100644
--- a/navix/components.py
+++ b/navix/components.py
@@ -19,48 +19,62 @@
 
 
 from __future__ import annotations
+from typing import Dict
 
 from jax import Array
 from flax import struct
 from jax.random import KeyArray
 import jax.numpy as jnp
 
+from .graphics import RenderingCache
+
+
+DISCARD_PILE_COORDS = jnp.asarray((0, -1), dtype=jnp.int32)
+DISCARD_PILE_IDX = jnp.asarray(-1, dtype=jnp.int32)
+EMPTY_POCKET_ID = jnp.asarray(-1, dtype=jnp.int32)
+
 
 class Component(struct.PyTreeNode):
     """A component is a part of the state of the environment."""
 
-    position: Array = jnp.zeros((1, 2), dtype=jnp.int32) - 1
-    """The (row, column) position of the entity in the grid, defaults to the discard pile (-1, -1)"""
-
 
 class Player(Component):
     """Players are entities that can act around the environment"""
 
+    position: Array = DISCARD_PILE_COORDS  # IntArray['b 2']
+    """The (row, column) position of the entity in the grid, defaults to the discard pile (-1, -1)"""
     # TODO(epignatelli): consider batching player over the number of players
     # to allow tranposing the entities pytree for faster computation
     # and to prepare the ground for multi-agent environments
-    tag: Array = jnp.asarray(1)
+    tag: Array = jnp.asarray(1)  # IntArray['2']
     """The tag of the component, used to identify the type of the component in `oobservations.categorical`"""
-    direction: Array = jnp.asarray(0)
+    # we mark direction as static because it is convenient for mapping observations (e.g. jnp.rot90(grid, k=direction)
+    # however, this is feasible because we only have 4 direcitions
+    # will it scale in multi-agent settings?
+    direction: Array = jnp.asarray(0, dtype=jnp.int32)  # IntArray['2']
     """The direction the entity: 0 = east, 1 = south, 2 = west, 3 = north"""
-    pocket: Array = jnp.asarray(0)
+    pocket: Array = EMPTY_POCKET_ID  # IntArray['2']
     """The id of the item in the pocket (0 if empty)"""
 
 
 class Goal(Component):
     """Goals are entities that can be reached by the player"""
 
-    tag: Array = jnp.ones((1,), dtype=jnp.int32) + 1
+    position: Array = DISCARD_PILE_COORDS[None]  # IntArray['b 2']
+    """The (row, column) position of the entity in the grid, defaults to the discard pile (-1, -1)"""
+    tag: Array = jnp.ones((1,), dtype=jnp.int32) + 1  # IntArray['b']
     """The tag of the component, used to identify the type of the component in `oobservations.categorical`"""
-    probability: Array = jnp.ones((1,), dtype=jnp.float32)
+    probability: Array = jnp.ones((1,), dtype=jnp.float32)  # FloatArray['b']
     """The probability of receiving the reward, if reached."""
 
 
-class Pickable(Component):
+class Key(Component):
     """Pickable items are world objects that can be picked up by the player.
     Examples of pickable items are keys, coins, etc."""
 
-    id: Array = jnp.zeros((1,), dtype=jnp.int32) - 1
+    position: Array = DISCARD_PILE_COORDS[None]  # IntArray['b 2']
+    """The (row, column) position of the entity in the grid, defaults to the discard pile (-1, -1)"""
+    id: Array = jnp.ones((1,), dtype=jnp.int32)  # IntArray['b']
     """The id of the item. If set, it must be >= 1."""
 
     @property
@@ -68,7 +82,7 @@ def tag(self):
         return -self.id
 
 
-class Consumable(Component):
+class Door(Component):
     """Consumable items are world objects that can be consumed by the player.
     Consuming an item requires a tool (e.g. a key to open a door).
     A tool is an id (int) of another item, specified in the `requires` field (-1 if no tool is required).
@@ -77,13 +91,15 @@ class Consumable(Component):
     Examples of consumables are doors (to open) food (to eat) and water (to drink), etc.
     """
 
-    requires: Array = jnp.zeros((1,), dtype=jnp.int32) - 1
+    position: Array = DISCARD_PILE_COORDS[None]  # IntArray['b 2']
+    """The (row, column) position of the entity in the grid, defaults to the discard pile (-1, -1)"""
+    requires: Array = EMPTY_POCKET_ID[None]  # IntArray['b']
     """The id of the item required to consume this item. If set, it must be >= 1."""
-    replacement: Array = jnp.zeros((1,), dtype=jnp.float32)
+    replacement: Array = jnp.zeros((1,), dtype=jnp.float32)  # IntArray['b']
     """The grid signature to replace the item with, usually 0 (floor). If set, it must be >= 1."""
 
     @property
-    def tag(self):
+    def tag(self) -> Array:  # -> IntArray['b']
         return self.requires
 
 
@@ -93,12 +109,47 @@ class State(struct.PyTreeNode):
     key: KeyArray
     """The random number generator state"""
     grid: Array
-    """The 2D-grid containing the ids of the entities in each position"""
+    """The base map of the environment that remains constant throughout the training"""
+    cache: RenderingCache
+    """The rendering cache to speed up rendering"""
     player: Player  # we can potentially extend this to multiple players easily
     """The player entity"""
     goals: Goal = Goal()
     """The goal entity, batched over the number of goals"""
-    keys: Pickable = Pickable()
+    keys: Key = Key()
     """The key entity, batched over the number of keys"""
-    doors: Consumable = Consumable()
+    doors: Door = Door()
     """The door entity, batched over the number of doors"""
+
+    def get_positions(self, axis: int = -1) -> Array:
+        return jnp.stack(
+            [
+                *self.keys.position,
+                *self.doors.position,
+                *self.goals.position,
+                self.player.position,
+            ],
+            axis=axis,
+        )
+
+    def get_tags(self, axis: int = -1) -> Array:
+        return jnp.stack(
+            [
+                *self.keys.tag,
+                *self.doors.tag,
+                *self.goals.tag,
+                self.player.tag,
+            ],
+            axis=axis,
+        )
+
+    def get_tiles(self, tiles_registry: Dict[str, Array], axis: int = 0) -> Array:
+        return jnp.stack(
+            [
+                *([tiles_registry["key"]] * len(self.keys.position)),
+                *([tiles_registry["door"]] * len(self.doors.position)),
+                *([tiles_registry["goal"]] * len(self.goals.position)),
+                tiles_registry["player"],
+            ],
+            axis=axis,
+        )
diff --git a/navix/environments/environment.py b/navix/environments/environment.py
index 6d943e1..8e626a7 100644
--- a/navix/environments/environment.py
+++ b/navix/environments/environment.py
@@ -22,18 +22,18 @@
 
 import abc
 from enum import IntEnum
-from typing import Any, Callable, Dict
+from typing import Any, Callable, Dict, Tuple
 import jax
 import jax.numpy as jnp
 from jax.random import KeyArray
 from jax import Array
 from flax import struct
 
-from .. import tasks
+
+from .. import tasks, graphics, terminations, observations
+from ..graphics import RenderingCache
 from ..components import State
-from .. import terminations
 from ..actions import ACTIONS
-from .. import observations
 
 
 class StepType(IntEnum):
@@ -68,7 +68,7 @@ class Environment(struct.PyTreeNode):
     max_steps: int = struct.field(pytree_node=False)
     gamma: float = struct.field(pytree_node=False, default=1.0)
     observation_fn: Callable[[State], Array] = struct.field(
-        pytree_node=False, default=observations.categorical
+        pytree_node=False, default=observations.none
     )
     reward_fn: Callable[[State, Array, State], Array] = struct.field(
         pytree_node=False, default=tasks.navigation
diff --git a/navix/environments/keydoor.py b/navix/environments/keydoor.py
index 310d74f..d2cc108 100644
--- a/navix/environments/keydoor.py
+++ b/navix/environments/keydoor.py
@@ -1,10 +1,12 @@
 import jax
 import jax.numpy as jnp
+from jax.random import KeyArray
 
-from navix.environments import Environment
-from navix.components import State, Player, Pickable, Consumable, Goal
-from navix.environments import Timestep
-from navix.grid import (
+from ..graphics import RenderingCache
+from ..environments import Environment
+from ..components import State, Player, Key, Door, Goal
+from ..environments import Timestep
+from ..grid import (
     two_rooms,
     random_positions,
     random_directions,
@@ -13,42 +15,32 @@
 
 
 class KeyDoor(Environment):
-    def reset(self, key) -> Timestep:
+    def reset(self, key: KeyArray) -> Timestep:
         key, k1, k2, k3, k4 = jax.random.split(key, 5)
 
         grid, wall_at = two_rooms(height=self.height, width=self.width, key=k4)
 
+        # add the door
+        door_pos = jnp.asarray([jax.random.randint(k3, (), 1, self.height - 1), wall_at])
+        doors = Door(position=door_pos[None], requires=jnp.asarray(3)[None])
+
         # spawn player and key in the first room
         out_of_bounds = jnp.asarray(self.height)
         first_room_mask = mask_by_coordinates(grid, (out_of_bounds, wall_at), jnp.less)
         first_room = jnp.where(first_room_mask, grid, -1)
-
         # player
         player_pos = random_positions(k1, first_room)
         player_dir = random_directions(k2)
         player = Player(position=player_pos, direction=player_dir)
-
         # key
         key_pos = random_positions(k2, first_room, exclude=player_pos[None])
-        keys = Pickable(position=key_pos[None], id=jnp.asarray(3)[None])
+        keys = Key(position=key_pos[None], id=jnp.asarray(3)[None])
 
         # spawn the goal in the second room
         second_room = jnp.where(first_room_mask, -1, grid)
-        goal_pos = random_positions(
-            k2, second_room, exclude=jnp.stack([player_pos, key_pos])
-        )
+        goal_pos = random_positions(k2, second_room)
         goals = Goal(position=goal_pos[None])
 
-        # add the door
-        door_coordinates = jnp.asarray(
-            [
-                jax.random.randint(k3, (), 1, self.height - 1),
-                wall_at,
-            ]
-        )
-        doors = Consumable(
-            position=door_coordinates[None], requires=jnp.asarray(3)[None]
-        )
 
         state = State(
             key=key,
@@ -57,6 +49,7 @@ def reset(self, key) -> Timestep:
             goals=goals,
             keys=keys,
             doors=doors,
+            cache=RenderingCache.init(grid),
         )
         return Timestep(
             t=jnp.asarray(0, dtype=jnp.int32),
diff --git a/navix/environments/room.py b/navix/environments/room.py
index 786552e..41be94b 100644
--- a/navix/environments/room.py
+++ b/navix/environments/room.py
@@ -19,14 +19,17 @@
 
 
 from __future__ import annotations
+from typing import Callable
 
 import jax
 import jax.numpy as jnp
 from jax.random import KeyArray
 
-from ..components import Goal, Player, State
+
+from ..components import Goal, Player
 from ..grid import random_positions, random_directions, room
-from .environment import Environment, Timestep
+from ..graphics import RenderingCache
+from .environment import Environment, Timestep, State
 
 
 class Room(Environment):
@@ -35,17 +38,24 @@ def reset(self, key: KeyArray) -> Timestep:
 
         # map
         grid = room(height=self.height, width=self.width)
-        positions = random_positions(k1, grid, n=2)
-        direction = random_directions(k2, n=1)
+        # TODO(epignatelli): if rendering gets slower, we can always
+        # split `reset`` into `init` and `reset`, start the cache in `init`
+        # and change it only when necessary in `reset`
+        # e.g., Room doesn't need to change the cache
+        # at every reset (and so many others) but KeyDoor does
+
         # player
-        player = Player(position=positions[0], direction=direction)
+        player_pos, goal_pos = random_positions(k1, grid, n=2)
+        direction = random_directions(k2, n=1)
+        player = Player(position=player_pos, direction=direction)
         # goal
-        goal = Goal(position=positions[1][None])
+        goal = Goal(position=goal_pos[None])
 
         # systems
         state = State(
             key=key,
             grid=grid,
+            cache=RenderingCache.init(grid),
             player=player,
             goals=goal,
         )
diff --git a/navix/graphics.py b/navix/graphics.py
index b2b0e3b..997974d 100644
--- a/navix/graphics.py
+++ b/navix/graphics.py
@@ -1,5 +1,10 @@
+from __future__ import annotations
+
+from typing import Dict, Tuple
+
 import jax
 import jax.numpy as jnp
+from flax import struct
 from jax import Array
 
 
@@ -15,6 +20,7 @@
 ORANGE = jnp.asarray([255, 128, 0], dtype=jnp.uint8)
 PINK = jnp.asarray([255, 0, 128], dtype=jnp.uint8)
 BROWN = jnp.asarray([128, 64, 0], dtype=jnp.uint8)
+GRAY_10 = jnp.asarray([230, 230, 230], dtype=jnp.uint8)
 GRAY_20 = jnp.asarray([205, 205, 205], dtype=jnp.uint8)
 GRAY_40 = jnp.asarray([153, 153, 153], dtype=jnp.uint8)
 GRAY_50 = jnp.asarray([128, 128, 128], dtype=jnp.uint8)
@@ -92,19 +98,40 @@ def colour_chart(size: int = TILE_SIZE) -> Array:
     return grid
 
 
-def colorise_tile(tile: Array, colour: Array, background: Array = GRAY_90) -> Array:
+class RenderingCache(struct.PyTreeNode):
+    patches: Array
+    """A flat set of patches representing the RGB values of each tile in the base map"""
+
+    @classmethod
+    def init(cls, grid: Array) -> RenderingCache:
+        background = render_background(grid)
+        patches = flatten_patches(background)
+
+        # add discard pile
+        patches = jnp.concatenate(
+            [
+                patches,
+                jnp.zeros((1, TILE_SIZE, TILE_SIZE, 3), dtype=jnp.uint8),
+            ],
+            axis=0,
+        )
+        return cls(patches=patches)
+
+
+def colorise_tile(tile: Array, colour: Array, background: Array = WHITE) -> Array:
+    assert tile.shape == (TILE_SIZE, TILE_SIZE), "Tile must be of size TILE_SIZE, TILE_SIZE, 3, got {}".format(tile.shape)
     tile = jnp.stack([tile] * colour.shape[0], axis=-1)
     tile = jnp.where(tile, colour, background)
     return tile
 
 
-def rectangle_tile(size: int = TILE_SIZE, colour: Array = BLACK) -> Array:
+def render_rectangle(size: int = TILE_SIZE, colour: Array = BLACK) -> Array:
     rectangle = jnp.ones((size - 2, size - 2), dtype=jnp.int32)
     rectangle = jnp.pad(rectangle, 1, "constant", constant_values=0)
     return colorise_tile(rectangle, colour)
 
 
-def triangle_east_tile(size: int = TILE_SIZE, colour: Array = RED) -> Array:
+def render_triangle_east(size: int = TILE_SIZE, colour: Array = RED) -> Array:
     triangle = jnp.ones((size, size), dtype=jnp.int32)
     triangle = jnp.tril(triangle, k=0)
     triangle = jnp.flip(triangle, axis=0)
@@ -113,25 +140,25 @@ def triangle_east_tile(size: int = TILE_SIZE, colour: Array = RED) -> Array:
     return colorise_tile(triangle, colour)
 
 
-def triangle_south_tile(size: int = TILE_SIZE, colour: Array = RED) -> Array:
-    triangle = triangle_east_tile(size)
+def render_triangle_south(size: int = TILE_SIZE, colour: Array = RED) -> Array:
+    triangle = render_triangle_east(size)
     triangle = jnp.rot90(triangle, k=3)
     return colorise_tile(triangle, colour)
 
 
-def triangle_west_tile(size: int = TILE_SIZE, colour: Array = RED) -> Array:
-    triangle = triangle_east_tile(size)
+def render_triangle_west(size: int = TILE_SIZE, colour: Array = RED) -> Array:
+    triangle = render_triangle_east(size)
     triangle = jnp.rot90(triangle, k=2)
     return colorise_tile(triangle, colour)
 
 
-def triangle_north_tile(size: int = TILE_SIZE, colour: Array = RED) -> Array:
-    triangle = triangle_east_tile(size)
+def render_triangle_north(size: int = TILE_SIZE, colour: Array = RED) -> Array:
+    triangle = render_triangle_east(size)
     triangle = jnp.rot90(triangle, k=1)
     return colorise_tile(triangle, colour)
 
 
-def diamond_tile(size: int = TILE_SIZE, colour: Array = GOLD) -> Array:
+def render_diamond(size: int = TILE_SIZE, colour: Array = GOLD) -> Array:
     diamond = jnp.ones((size, size), dtype=jnp.int32)
     diamond = jnp.tril(diamond, k=size // 2.5)
     diamond = jnp.flip(diamond, axis=0)
@@ -143,7 +170,7 @@ def diamond_tile(size: int = TILE_SIZE, colour: Array = GOLD) -> Array:
     return colorise_tile(diamond, colour)
 
 
-def door_tile(size: int = TILE_SIZE, colour: Array = BROWN) -> Array:
+def render_door(size: int = TILE_SIZE, colour: Array = BROWN) -> Array:
     frame_size = TILE_SIZE - 6
     door = jnp.zeros((frame_size, frame_size), dtype=jnp.int32)
     door = jnp.pad(door, 1, "constant", constant_values=1)
@@ -159,7 +186,7 @@ def door_tile(size: int = TILE_SIZE, colour: Array = BROWN) -> Array:
     return colorise_tile(door, colour)
 
 
-def key_tile(size: int = TILE_SIZE, colour: Array = BRONZE) -> Array:
+def render_key(size: int = TILE_SIZE, colour: Array = BRONZE) -> Array:
     key = jnp.zeros((size, size), dtype=jnp.int32)
 
     # Handle (Round Part)
@@ -195,17 +222,81 @@ def key_tile(size: int = TILE_SIZE, colour: Array = BRONZE) -> Array:
     return colorise_tile(key, colour)
 
 
-def floor_tile(size: int = TILE_SIZE, colour: Array = GRAY_90) -> Array:
+def render_floor(size: int = TILE_SIZE, colour: Array = WHITE) -> Array:
     floor = jnp.ones((size - 2, size - 2), dtype=jnp.int32)
     floor = jnp.pad(floor, 1, "constant", constant_values=0)
-    return colorise_tile(floor, colour, background=GRAY_50)
+    return colorise_tile(floor, colour, background=GRAY_10)
 
 
-def wall_tile(size: int = TILE_SIZE, colour: Array = GRAY_50) -> Array:
+def render_wall(size: int = TILE_SIZE, colour: Array = GRAY_80) -> Array:
     wall = jnp.ones((size, size), dtype=jnp.int32)
     return colorise_tile(wall, colour)
 
 
-def mosaic(grid: Array, tile: Array) -> Array:
+def tile_grid(grid: Array, tile: Array) -> Array:
     tiled = jnp.tile(tile, (*grid.shape, 1))
     return jnp.asarray(tiled, dtype=jnp.uint8)
+
+
+TILES_REGISTRY: Dict[str, Array] = {
+    "wall": render_wall(),
+    "floor": render_floor(),
+    "player": render_triangle_east(),
+    "goal": render_diamond(),
+    "key": render_key(),
+    "door": render_door(),
+}
+
+
+def render_background(
+    grid: Array, tiles_registry: Dict[str, Array] = TILES_REGISTRY
+) -> Array:
+    image_width = grid.shape[0] * TILE_SIZE
+    image_height = grid.shape[1] * TILE_SIZE
+    n_channels = 3
+
+    background = jnp.zeros((image_height, image_width, n_channels), dtype=jnp.uint8)
+    grid_resized = jax.image.resize(
+        grid, (grid.shape[0] * TILE_SIZE, grid.shape[1] * TILE_SIZE), method="nearest"
+    )
+
+    mask = jnp.asarray(grid_resized, dtype=bool)  # 0 = floor, 1 = wall
+    floor_tile = tile_grid(grid, tiles_registry["floor"])
+    wall_tile = tile_grid(grid, tiles_registry["wall"])
+    background = jnp.where(mask[..., None], wall_tile, floor_tile)
+    return background
+
+
+def flatten_patches(image: Array, patch_size: Tuple[int, int] = (TILE_SIZE, TILE_SIZE)) -> Array:
+    height = image.shape[0] // patch_size[0]
+    width = image.shape[1] // patch_size[1]
+    n_channels = image.shape[2]
+
+    grid = image.reshape(height, patch_size[0], width, patch_size[1], n_channels)
+
+    # Swap the first and second axes of the grid to revert the stacking order
+    grid = jnp.swapaxes(grid, 1, 2)
+
+    # Reshape the grid of tiles into the original list of tiles
+    patches = grid.reshape(height * width, patch_size[0], patch_size[1], n_channels)
+
+    return patches
+
+
+def unflatten_patches(patches: Array, image_size: Tuple[int, int]) -> Array:
+    image_height = image_size[0]
+    image_width = image_size[1]
+    patch_height = patches.shape[1]
+    patch_width = patches.shape[2]
+    n_channels = patches.shape[3]
+
+    # Reshape the list of tiles into a 2D grid
+    grid = patches.reshape(image_height // patch_height, image_width // patch_width, patch_height, patch_width, n_channels)
+
+    # Swap the first and second axes of the grid to change the order of stacking
+    grid = jnp.swapaxes(grid, 1, 2)
+
+    # Reshape and stack the grid tiles horizontally and vertically to form the final image
+    image = grid.reshape(image_height, image_width, n_channels)
+
+    return image
diff --git a/navix/grid.py b/navix/grid.py
index bd2b772..26c7378 100644
--- a/navix/grid.py
+++ b/navix/grid.py
@@ -31,17 +31,23 @@
 Coordinates = Tuple[Array, Array]
 
 
+def coordinates(grid: Array) -> Coordinates:
+    return tuple(jnp.mgrid[0 : grid.shape[0], 0 : grid.shape[1]])
+
+
 def idx_from_coordinates(grid: Array, coordinates: Array):
-    """Converts a 2D coordinate (col, row) into a flat index"""
+    """Converts a batch of 2D coordinates [(col, row), ...] into a flat index"""
+    coordinates = coordinates.T
+    assert coordinates.shape[0] == 2, coordinates.shape
+
     idx = coordinates[0] * grid.shape[1] + coordinates[1]
     return jnp.asarray(idx, dtype=jnp.int32)
 
 
 def coordinates_from_idx(grid: Array, idx: Array):
     """Converts a flat index into a 2D coordinate (col, row)"""
-    col, row = jnp.divmod(idx, grid.shape[1])
-    coords = jnp.stack([col, row])
-    return jnp.asarray(coords, dtype=jnp.int32)
+    coords = jnp.divmod(idx, grid.shape[1])
+    return jnp.asarray(coords, dtype=jnp.int32).T
 
 
 def mask_by_coordinates(
@@ -49,9 +55,12 @@ def mask_by_coordinates(
     address: Coordinates,
     comparison_fn: Callable[[Array, Array], Array] = jnp.greater_equal,
 ) -> Array:
-    """Returns a mask of the same shape as `grid` where the value is 1 if the
+    """This is a workaround to compute dynamicly-sized masks in XLA,
+    which would not be possible otherwise.
+    Returns a mask of the same shape as `grid` where the value is 1 if the
     corresponding element in `grid` satisfies the `comparison_fn` with the
-    corresponding element in `address` (col, row) and 0 otherwise."""
+    corresponding element in `address` (col, row) and 0 otherwise.
+    """
     mesh = jnp.mgrid[0 : grid.shape[0], 0 : grid.shape[1]]
     cond_1 = comparison_fn(mesh[0], address[0])
     cond_2 = comparison_fn(mesh[1], address[1])
@@ -59,6 +68,49 @@ def mask_by_coordinates(
     return mask
 
 
+def translate(
+    position: Array, direction: Array, modulus: Array = jnp.asarray(1)
+) -> Array:
+    moves = (
+        lambda position: position + jnp.asarray((0, modulus)),  # east
+        lambda position: position + jnp.asarray((modulus, 0)),  # south
+        lambda position: position + jnp.asarray((0, -modulus)),  # west
+        lambda position: position + jnp.asarray((-modulus, 0)),  # north
+    )
+    return jax.lax.switch(direction, moves, position)
+
+
+def translate_forward(position: Array, forward_direction: Array, modulus: Array):
+    return translate(position, forward_direction, modulus)
+
+
+def translate_left(position: Array, forward_direction: Array, modulus: Array):
+    return translate(position, (forward_direction + 3) % 4, modulus)
+
+
+def translate_right(position: Array, forward_direction: Array, modulus: Array):
+    return translate(position, (forward_direction + 1) % 4, modulus)
+
+
+def rotate(direction: Array, spin: int) -> Array:
+    return (direction + spin) % 4
+
+
+def random_positions(
+    key: KeyArray, grid: Array, n: int = 1, exclude: Array = jnp.asarray((-1, -1))
+) -> Array:
+    probs = grid.reshape(-1)
+    indices = idx_from_coordinates(grid, exclude)
+    probs = probs.at[indices].set(-1) + 1.0
+    idx = jax.random.categorical(key, jnp.log(probs), shape=(n,))
+    position = coordinates_from_idx(grid, idx)
+    return position.squeeze()
+
+
+def random_directions(key: KeyArray, n=1) -> Array:
+    return jax.random.randint(key, (n,), 0, 4).squeeze()
+
+
 def room(height: int, width: int):
     """A grid of ids of size `width` x `height`"""
     grid = jnp.zeros((height, width), dtype=jnp.int32)
@@ -77,30 +129,9 @@ def two_rooms(height: int, width: int, key: KeyArray) -> Tuple[Array, Array]:
     return grid, wall_at
 
 
-def random_positions(
-    key: KeyArray, grid: Array, n=1, exclude: Array = jnp.asarray([(-1, -1)])
-) -> Array:
-    # all floor tiles
-    # TODO(epignatelli): switch to all walkable tiles
-    mask = jnp.where(grid, 0, 1).reshape((-1,))  # all floor tiles
-
-    # temporarily set excluded positions to 0 probability
-    mask = jnp.min(
-        jax.vmap(lambda position: mask.at[idx_from_coordinates(grid, position)].set(0))(
-            exclude
-        ),
-        axis=0,
-    )
-
-    log_probs = jnp.log(mask)
-    idx = jax.random.categorical(key, log_probs, shape=(n,))
-
-    positions = jax.vmap(coordinates_from_idx, in_axes=(None, 0))(grid, idx)
-    return positions.T.squeeze()
-
-
-def random_directions(key: KeyArray, n=1) -> Array:
-    return jax.random.randint(key, (n,), 0, 4).squeeze()
+def merge_maps(grid_1: Array, grid_2: Array) -> Array:
+    """Merge two overlaying maps of floors (0) and walls (-1) into a single grid"""
+    return jnp.asarray(grid_1 + grid_2, dtype=jnp.int32)
 
 
 def from_ascii_map(ascii_map: str, mapping: Dict[str, int] = {}) -> Array:
diff --git a/navix/observations.py b/navix/observations.py
index 43df205..2af05fa 100644
--- a/navix/observations.py
+++ b/navix/observations.py
@@ -19,108 +19,85 @@
 
 
 from __future__ import annotations
+from typing import Dict, Tuple
 
-from jax import Array
 import jax
 import jax.numpy as jnp
+from jax import Array
 
-from .components import State
-from .graphics import (
-    triangle_east_tile,
-    diamond_tile,
-    door_tile,
-    key_tile,
-    floor_tile,
-    wall_tile,
-    mosaic,
-    TILE_SIZE,
-)
-
-
-def third_person_view(state: State, radius: int) -> Array:
-    raise NotImplementedError()
-
-
-def first_person_view(state: State, radius: int) -> Array:
-    raise NotImplementedError()
-
-
-def categorical(state: State) -> Array:
-    # updates are in reverse order of display
-    # place keys (keys are represented with opposite sign of the door they open)
-    grid = jnp.max(
-        jax.vmap(lambda key: state.grid.at[tuple(key.position)].set(-key.tag))(
-            state.keys
-        ),
-        axis=0,
-    )
-    # place doors
-    grid = jnp.max(
-        jax.vmap(lambda door: grid.at[tuple(door.position)].set(door.tag))(state.doors),
-        axis=0,
-    )
-    # place goals
-    grid = jnp.max(
-        jax.vmap(lambda goal: grid.at[tuple(goal.position)].set(goal.tag))(state.goals),
-        axis=0,
-    )
-    # place player last, always on top
-    grid = grid.at[tuple(state.player.position)].set(state.player.tag)
-    return grid
-
-
-def rgb(state: State) -> Array:
-    positions = jnp.stack(
-        [
-            *state.keys.position,
-            *state.doors.position,
-            *state.goals.position,
-            state.player.position,
-        ]
-    )
-
-    tiles = jnp.stack(
-        [
-            *([key_tile()] * len(state.keys.position)),
-            *([door_tile()] * len(state.doors.position)),
-            *([diamond_tile()] * len(state.goals.position)),
-            triangle_east_tile(),
-        ]
-    )
-
-    def draw(carry, x):
-        image = carry
-        mask, tile = x
-        mask = jax.image.resize(
-            mask,
-            (mask.shape[0] * TILE_SIZE, mask.shape[1] * TILE_SIZE),
-            method="nearest",
-        )
-        mask = jnp.stack([mask] * tile.shape[-1], axis=-1)
-        tiled = mosaic(state.grid, tile)
-        image = jnp.where(mask, tiled, image)
-        return (image), ()
-
-    def body_fun(carry, x):
-        position, tile = x
-        mask = jnp.zeros_like(state.grid).at[tuple(position)].set(1)
-        return draw(carry, (mask, tile))
-
-    background = jnp.zeros(
-        (state.grid.shape[0] * TILE_SIZE, state.grid.shape[1] * TILE_SIZE, 3),
-        dtype=jnp.uint8,
-    )
-
-    # add floor
-    floor_mask = jnp.where(state.grid == 0, 1, 0)
-    floor = floor_tile()
-    background, _ = draw(background, (floor_mask, floor))
-
-    # add walls
-    wall_mask = jnp.where(state.grid == -1, 1, 0)
-    wall = wall_tile()
-    background, _ = draw(background, (wall_mask, wall))
-
-    # add entities
-    image, _ = jax.lax.scan(body_fun, background, (positions, tiles))
-    return image  # type: ignore
+from . import graphics
+from .components import DISCARD_PILE_IDX, State
+from .grid import idx_from_coordinates
+
+
+# def _view_extremes(state: State, radius: Array) -> Tuple[Array, Array]:
+#     pos_left = translate_left(state.player.position, state.player.direction, radius)
+#     pos_forward_left = translate_forward(pos_left, state.player.direction, radius)
+#     pos_right = translate_right(state.player.position, state.player.direction, radius)
+#     pos_forward_right = translate_forward(pos_right, state.player.direction, radius)
+
+#     all_pos = jnp.stack([pos_left, pos_forward_left, pos_right, pos_forward_right], axis=0)
+
+#     north_west = jnp.min(all_pos, axis=0)
+#     south_east = jnp.max(all_pos, axis=0)
+
+#     return north_west, south_east
+
+
+# def _view_mask(state: State, radius: Array) -> Array:
+#     north_west, south_east = _view_extremes(state, radius)
+#     mask = state.grid.at[north_west[0]:south_east[0] + 1, north_west[1]:south_east[1] + 1].set(1)
+#     return mask
+
+
+# def _first_person_crop(state: State, radius: Array) -> Array:
+#     north_west, south_east = _view_extremes(state, radius)
+#     view = state.grid[north_west[0]:south_east[0] + 1, north_west[1]:south_east[1] + 1]
+#     view = jnp.rot90(view, k=state.player.direction + 1)
+#     return view
+
+
+def none(
+    state: State,
+    # cache: graphics.RenderingCache,
+    tiles_registry: Dict[str, Array] = graphics.TILES_REGISTRY,
+) -> Array:
+    return jnp.asarray(())
+
+
+def categorical(
+    state: State,
+    # cache: graphics.RenderingCache,
+    tiles_registry: Dict[str, Array] = graphics.TILES_REGISTRY,
+) -> Array:
+    # get idx of entity on the set of patches
+    indices = idx_from_coordinates(state.grid, state.get_positions(axis=0))
+    # get tags corresponding to the entities
+    tags = state.get_tags(axis=0)
+    # set tags on the flat set of patches
+    shape = state.grid.shape
+    grid = state.grid.reshape(-1).at[indices].set(tags)
+    # unflatten patches to reconstruct the grid
+    return grid.reshape(shape)
+
+
+def rgb(
+    state: State,
+    # cache: graphics.RenderingCache,
+    tiles_registry: Dict[str, Array] = graphics.TILES_REGISTRY,
+) -> Array:
+    # for 1-d vs 2-d indexing benchamarks
+    # see https://github.com/epignatelli/navix/tree/observation/2dindexing
+
+    # get idx of entity on the flat set of patches
+    indices = idx_from_coordinates(state.grid, state.get_positions(axis=0))
+    # get tiles corresponding to the entities
+    tiles = state.get_tiles(tiles_registry, axis=0)
+    # set tiles on the flat set of patches
+    patches = state.cache.patches.at[indices].set(tiles)
+    # remove discard pile
+    patches = patches[:DISCARD_PILE_IDX]
+    # unflatten patches to reconstruct the image
+    image_size = (state.grid.shape[0] * graphics.TILE_SIZE, state.grid.shape[1] * graphics.TILE_SIZE)
+    image = graphics.unflatten_patches(patches, image_size)
+    return image


Test patch:
diff --git a/tests/performance/grid.py b/tests/performance/grid.py
new file mode 100644
index 0000000..9056969
--- /dev/null
+++ b/tests/performance/grid.py
@@ -0,0 +1,52 @@
+import time
+from timeit import repeat
+
+import jax
+import jax.numpy as jnp
+import navix as nx
+
+
+N_TIMEIT_LOOPS = 10
+N_REPEAT = 100
+N_TIMESTEPS = 1000
+N_SEEDS = 100
+
+
+def test_observation():
+    def test(seed):
+        env = nx.environments.Room.create(
+            height=10, width=5, max_steps=100, observation_fn=nx.observations.none
+        )
+        key = jax.random.PRNGKey(seed)
+        timestep = env.reset(key)
+
+        actions = jax.random.randint(key, (100,), 0, 6)
+        timestep = jax.lax.scan(lambda c, x: (env.step(c, x), ()), timestep, actions)[0]
+        return timestep
+
+    # profile navix scanned
+    print(
+        "Profiling, N_SEEDS = {}, N_TIMESTEPS = {}".format(
+            N_SEEDS, N_TIMESTEPS
+        )
+    )
+
+    seeds = jnp.arange(N_SEEDS)
+
+    print(f"\tCompiling {test}...")
+    start = time.time()
+    test_jit = jax.jit(jax.vmap(test)).lower(seeds).compile()
+    print("\tCompiled in {:.2f}s".format(time.time() - start))
+
+    print("\tRunning ...")
+    res = repeat(
+        lambda: test_jit(seeds).observation.block_until_ready(),
+        number=N_TIMEIT_LOOPS,
+        repeat=N_REPEAT,
+    )
+    res = jnp.asarray(res)
+    print(f"\t {jnp.mean(res)} ± {jnp.std(res)}")
+
+
+if __name__ == "__main__":
+    test_observation()
diff --git a/tests/performance/minigrid.py b/tests/performance/minigrid.py
new file mode 100644
index 0000000..1138591
--- /dev/null
+++ b/tests/performance/minigrid.py
@@ -0,0 +1,74 @@
+import jax
+import jax.numpy as jnp
+import navix as nx
+
+import gymnasium
+import random
+import time
+
+from timeit import repeat
+
+N_TIMEIT_LOOPS = 3
+N_REPEAT = 5
+N_TIMESTEPS = 1000
+N_SEEDS = 10_000
+
+
+def profile_navix_scan(seed):
+    env = nx.environments.Room.create(
+        height=10, width=5, max_steps=100, observation_fn=nx.observations.categorical
+    )
+    key = jax.random.PRNGKey(4)
+    timestep = env.reset(key)
+    actions = jax.random.randint(key, (N_TIMESTEPS,), 0, 6)
+
+    timestep = jax.lax.scan(
+        lambda carry, x: (env.step(carry, x), ()), timestep, actions
+    )[0]
+
+    return timestep
+
+
+def profile_minigrid(seed):
+    env = gymnasium.make("MiniGrid-Empty-16x16-v0", render_mode=None)
+    observation, info = env.reset(seed=42)
+    for _ in range(N_TIMESTEPS):
+        action = random.randint(0, 4)
+        observation, reward, terminated, truncated, info = env.step(action)
+
+        if terminated or truncated:
+            observation, info = env.reset()
+    env.close()
+    return observation
+
+
+if __name__ == "__main__":
+    # profile navix scanned
+    print(
+        "Profiling navix with `scan`, N_SEEDS = {}, N_TIMESTEPS = {}".format(
+            N_SEEDS, N_TIMESTEPS
+        )
+    )
+    seeds = jnp.arange(N_SEEDS)
+
+    print(f"\tCompiling {profile_navix_scan}...")
+    start = time.time()
+    f_scan = jax.jit(jax.vmap(profile_navix_scan)).lower(seeds).compile()
+    print("\tCompiled in {:.2f}s".format(time.time() - start))
+
+    print("\tRunning ...")
+    res_navix = repeat(
+        lambda: f_scan(seeds).observation.block_until_ready(),
+        number=N_TIMEIT_LOOPS,
+        repeat=N_REPEAT,
+    )
+    res_navix = jnp.asarray(res_navix)
+    print(f"\t {jnp.mean(res_navix)} ± {jnp.std(res_navix)}")
+
+    # profile minigrid
+    print("Profiling minigrid, N_SEEDS = 1, N_TIMESTEPS = {}".format(N_TIMESTEPS))
+    res_minigrid = repeat(
+        lambda: profile_minigrid(0), number=N_TIMEIT_LOOPS, repeat=N_REPEAT
+    )
+    res_minigrid = jnp.asarray(res_minigrid)
+    print(f"\t {jnp.mean(res_minigrid)} ± {jnp.std(res_minigrid)}")
diff --git a/tests/performance/minigrid_report.txt b/tests/performance/minigrid_report.txt
new file mode 100644
index 0000000..63ec334
--- /dev/null
+++ b/tests/performance/minigrid_report.txt
@@ -0,0 +1,6 @@
+Profiling navix with `scan`, N_SEEDS = 10000, N_TIMESTEPS = 1000
+	Compiling <function profile_navix_scan at 0x7f1318bcd9d0>...
+	Compiled in 0.78s
+	Running ...
+	 0.1666584610939026 ± 0.021019447594881058
+Profiling minigrid, N_SEEDS = 1, N_TIMESTEPS = 1000
diff --git a/tests/performance/observations.py b/tests/performance/observations.py
new file mode 100644
index 0000000..2379880
--- /dev/null
+++ b/tests/performance/observations.py
@@ -0,0 +1,54 @@
+import time
+from timeit import repeat
+
+import jax
+import jax.numpy as jnp
+import navix as nx
+
+
+N_TIMEIT_LOOPS = 10
+N_REPEAT = 30
+N_TIMESTEPS = 1000
+N_SEEDS = 10
+
+
+def test_observation(observation_fn):
+    def test(seed):
+        env = nx.environments.KeyDoor(
+            height=10, width=5, max_steps=100, gamma=1.0, observation_fn=observation_fn
+        )
+        key = jax.random.PRNGKey(seed)
+        timestep = env.reset(key)
+
+        actions = jax.random.randint(key, (100,), 0, 6)
+        timestep = jax.lax.scan(lambda c, x: (env.step(c, x), ()), timestep, actions)[0]
+        return timestep
+
+    # profile navix scanned
+    print(
+        "Profiling observation {}, N_SEEDS = {}, N_TIMESTEPS = {}".format(
+            observation_fn, N_SEEDS, N_TIMESTEPS
+        )
+    )
+
+    seeds = jnp.arange(N_SEEDS)
+
+    print(f"\tCompiling {observation_fn}...")
+    start = time.time()
+    test_jit = jax.jit(jax.vmap(test)).lower(seeds).compile()
+    print("\tCompiled in {:.2f}s".format(time.time() - start))
+
+    print("\tRunning ...")
+    res = repeat(
+        lambda: test_jit(seeds).observation.block_until_ready(),
+        number=N_TIMEIT_LOOPS,
+        repeat=N_REPEAT,
+    )
+    res = jnp.asarray(res)
+    print(f"\t {jnp.mean(res)} ± {jnp.std(res)}")
+
+
+if __name__ == "__main__":
+    test_observation(nx.observations.none)
+    test_observation(nx.observations.categorical)
+    test_observation(nx.observations.rgb)
diff --git a/tests/performance/observations_keydoor_report.txt b/tests/performance/observations_keydoor_report.txt
new file mode 100644
index 0000000..19244ed
--- /dev/null
+++ b/tests/performance/observations_keydoor_report.txt
@@ -0,0 +1,15 @@
+Profiling observation <function none at 0x7fa94c162940>, N_SEEDS = 10, N_TIMESTEPS = 1000
+	Compiling <function none at 0x7fa94c162940>...
+	Compiled in 1.35s
+	Running ...
+	 0.07763808965682983 ± 0.004634591285139322
+Profiling observation <function categorical at 0x7fa94c162a60>, N_SEEDS = 10, N_TIMESTEPS = 1000
+	Compiling <function categorical at 0x7fa94c162a60>...
+	Compiled in 1.33s
+	Running ...
+	 0.09059003740549088 ± 0.0004746166814584285
+Profiling observation <function rgb at 0x7fa94c1629d0>, N_SEEDS = 10, N_TIMESTEPS = 1000
+	Compiling <function rgb at 0x7fa94c1629d0>...
+	Compiled in 1.47s
+	Running ...
+	 0.14711560308933258 ± 0.0005664637428708375
diff --git a/tests/performance/observations_room_report.txt b/tests/performance/observations_room_report.txt
new file mode 100644
index 0000000..c0992a6
--- /dev/null
+++ b/tests/performance/observations_room_report.txt
@@ -0,0 +1,15 @@
+Profiling observation <function none at 0x7ff30c055940>, N_SEEDS = 10, N_TIMESTEPS = 1000
+	Compiling <function none at 0x7ff30c055940>...
+	Compiled in 0.99s
+	Running ...
+	 0.04586170241236687 ± 0.00395115977153182
+Profiling observation <function categorical at 0x7ff30c055a60>, N_SEEDS = 10, N_TIMESTEPS = 1000
+	Compiling <function categorical at 0x7ff30c055a60>...
+	Compiled in 0.99s
+	Running ...
+	 0.058218248188495636 ± 0.0005350122228264809
+Profiling observation <function rgb at 0x7ff30c0559d0>, N_SEEDS = 10, N_TIMESTEPS = 1000
+	Compiling <function rgb at 0x7ff30c0559d0>...
+	Compiled in 1.08s
+	Running ...
+	 0.1171514168381691 ± 0.00043898896547034383
diff --git a/tests/performance/profiling.py b/tests/performance/profiling.py
new file mode 100644
index 0000000..6b288a7
--- /dev/null
+++ b/tests/performance/profiling.py
@@ -0,0 +1,51 @@
+import jax
+import jax.numpy as jnp
+import navix as nx
+
+
+N_TIMESTEPS = 1000
+
+
+def f():
+    seed = 0
+    env = nx.environments.Room(16, 16, 8, observation_fn=nx.observations.rgb)
+    key = jax.random.PRNGKey(seed)
+    timestep = env.reset(key)
+    timestep = env.step(timestep, jnp.asarray(3))
+    timestep = env.step(timestep, jnp.asarray(3))
+    timestep = env.step(timestep, jnp.asarray(3))
+    timestep = env.step(timestep, jnp.asarray(3))
+    timestep = env.step(timestep, jnp.asarray(3))
+    timestep = env.step(timestep, jnp.asarray(3))
+    timestep = env.step(timestep, jnp.asarray(3))
+    return timestep
+
+
+def f_scan():
+    seed = 0
+    env = nx.environments.Room(16, 16, 8, observation_fn=nx.observations.rgb)
+    key = jax.random.PRNGKey(seed)
+    timestep = env.reset(key)
+
+    def body_fun(carry, x):
+        timestep = carry
+        action = x
+        timestep = env.step(timestep, jnp.asarray(action))
+        return timestep, ()
+
+    actions = jax.random.randint(key, (N_TIMESTEPS,), 0, 6)
+    timestep = jax.lax.scan(
+        lambda c, x: (env.step(c, x), ()),
+        timestep,
+        jnp.asarray(actions, dtype=jnp.int32),
+    )[0]
+    return timestep
+
+
+f_jit = jax.jit(f)
+f_scan_jit = jax.jit(f_scan)
+
+
+with jax.profiler.trace("/tmp/jax-trace", create_perfetto_link=True):
+    timestep = f_scan_jit()
+    timestep.observation.block_until_ready()
diff --git a/tests/test_actions.py b/tests/test_actions.py
index d33a5f4..b342251 100644
--- a/tests/test_actions.py
+++ b/tests/test_actions.py
@@ -2,15 +2,22 @@
 import jax.numpy as jnp
 
 import navix as nx
+from navix.components import EMPTY_POCKET_ID, DISCARD_PILE_COORDS
 
 
 def test_rotation():
     direction = jnp.asarray(0)
 
+    key = jax.random.PRNGKey(0)
+    grid = jnp.zeros((3, 3), dtype=jnp.int32)
+    player = nx.components.Player(position=jnp.asarray((1, 1)), direction=direction)
+    cache = nx.graphics.RenderingCache(grid)
+
     state = nx.components.State(
-        jax.random.PRNGKey(0),
-        grid=jnp.zeros((3, 3), dtype=jnp.int32),
-        player=nx.components.Player(position=jnp.asarray((1, 1)), direction=direction),
+        grid=grid,
+        player=player,
+        cache=cache,
+        key=key,
     )
 
     msg = "Expected direction to be {}, got {}"
@@ -37,38 +44,178 @@ def test_rotation():
 
 
 def test_walkable():
-    height = 6
-    width = 18
-    env = nx.environments.KeyDoor(
-        height=height,
-        width=width,
-        max_steps=100,
-        observation_fn=nx.observations.categorical,
+    heigh, width = 5, 5
+    grid = jnp.zeros((heigh - 2, width - 2), dtype=jnp.int32)
+    grid = jnp.pad(grid, pad_width=1, mode="constant", constant_values=1)
+    key = jax.random.PRNGKey(0)
+    player = nx.components.Player(position=jnp.asarray((1, 1)), direction=jnp.asarray(0))
+    goals = nx.components.Goal(position=jnp.asarray((3, 3))[None])
+    keys = nx.components.Key(position=jnp.asarray((3, 1))[None])
+    doors = nx.components.Door(position=jnp.asarray((1, 3))[None])
+    cache = nx.graphics.RenderingCache(grid)
+    # Looks like this
+    # -1 -1 -1 -1 -1
+    # -1  P  0  D -1
+    # -1  0  0  0 -1
+    # -1  K  0  G -1
+    # -1 -1 -1 -1 -1
+
+    state = nx.components.State(
+        grid=grid,
+        player=player,
+        goals=goals,
+        keys=keys,
+        doors=doors,
+        cache=cache,
+        key=key,
     )
 
+    def check_forward_position(state):
+        prev_pos = state.player.position
+        pos = nx.grid.translate_forward(prev_pos, state.player.direction, jnp.asarray(1))
+        assert not nx.actions._walkable(state, pos), (
+            "Expected position {} to be not walkable, since it is a {}".format(pos, state.grid[tuple(pos)])
+        )
+
+    state = nx.actions.forward(state)
+
+    # should not be able to walk on/through a door
+    check_forward_position(state)
+
+    state = nx.actions.backward(state)
+    state = nx.actions.rotate_cw(state)
+    state = nx.actions.forward(state)
+
+    # should not be able to walk on/through a key
+    check_forward_position(state)
+
+    state = nx.actions.rotate_ccw(state)
+    state = nx.actions.forward(state)
+    state = nx.actions.rotate_cw(state)
+    state = nx.actions.forward(state)
+
+    # should not be able to walk on/through a wall
+    check_forward_position(state)
+
+    state = nx.actions.rotate_ccw(state)
+
+    # should be able to walk on/through a goal
+    state = nx.actions.forward(state)
+    # but not on/through a wall
+    check_forward_position(state)
+
+
+def test_pickup():
+    heigh, width = 5, 5
+    grid = jnp.zeros((heigh - 2, width - 2), dtype=jnp.int32)
+    grid = jnp.pad(grid, pad_width=1, mode="constant", constant_values=1)
     key = jax.random.PRNGKey(0)
-    timestep = env.reset(key)
-    actions = (
-        2,
-        3, # in front of key after this
-    )
-    actions_stuck = (
-        3, # should not be able to move forward
-        3, # should not be able to move forward
-        2, # rotate towards the wall
-        3, # should not be able to move forward
-        3, # should not be able to move forward
-    )
-    for action in actions:
-        timestep = env.step(timestep, jnp.asarray(action))
-        print(timestep.state.player.position)
-
-    for action in actions_stuck:
-        next_timestep = env.step(timestep, jnp.asarray(action))
-        print(timestep.state.player.position)
-        assert jnp.array_equal(timestep.state.player.position, next_timestep.state.player.position)
-        timestep = next_timestep
+    player = nx.components.Player(position=jnp.asarray((1, 1)), direction=jnp.asarray(1))
+    goals = nx.components.Goal(position=jnp.asarray((3, 3))[None])
+    keys = nx.components.Key(position=jnp.asarray((2, 1))[None], id=jnp.asarray(1)[None])
+    doors = nx.components.Door(position=jnp.asarray((1, 3))[None], requires=jnp.asarray(1))
+    cache = nx.graphics.RenderingCache(grid)
+
+    # Looks like this
+    """
+    #  #  #  #  #
+    #  P  .  D  #
+    #  K  .  .  #
+    #  .  .  G  #
+    #  #  #  #  #
+    """
+    state = nx.components.State(
+        grid=grid,
+        player=player,
+        goals=goals,
+        keys=keys,
+        doors=doors,
+        cache=cache,
+        key=key,
+    )
+
+    # check that the player has no keys
+    assert jnp.array_equal(state.player.pocket, EMPTY_POCKET_ID), (
+        "Expected player to have pocket {}, got {}".format(EMPTY_POCKET_ID, state.player.pocket)
+    )
+
+    # pick up the key
+    state = nx.actions.pickup(state)
+
+    # check that the player has the key
+    expected_pocket = jnp.asarray(1)
+    assert jnp.array_equal(state.player.pocket, expected_pocket), (
+        "Expected player to have key {}, got {}".format(expected_pocket, state.player.pocket)
+    )
+
+    # check that the key is no longer on the grid
+    assert jnp.array_equal(state.keys.position[0], DISCARD_PILE_COORDS), (
+        "Expected key to be at {}, got {}".format(DISCARD_PILE_COORDS, state.keys.position)
+    )
+
+
+def test_open():
+    heigh, width = 5, 5
+    grid = jnp.zeros((heigh - 2, width - 2), dtype=jnp.int32)
+    grid = jnp.pad(grid, pad_width=1, mode="constant", constant_values=1)
+    key = jax.random.PRNGKey(0)
+    player = nx.components.Player(position=jnp.asarray((1, 1)), direction=jnp.asarray(0))
+    goals = nx.components.Goal(position=jnp.asarray((3, 3))[None])
+    keys = nx.components.Key(position=jnp.asarray((3, 1))[None], id=jnp.asarray(1))
+    doors = nx.components.Door(position=jnp.asarray((1, 3))[None], requires=jnp.asarray(1))
+    cache = nx.graphics.RenderingCache(grid)
+
+    # Looks like this
+    # W  W  W  W  W
+    # W  P  0  D  W
+    # W  0  0  0  W
+    # W  K  0  G  W
+    # W  W  W  W  W
+
+    state = nx.components.State(
+        grid=grid,
+        player=player,
+        goals=goals,
+        keys=keys,
+        doors=doors,
+        cache=cache,
+        key=key,
+    )
+
+    # check that the player has no keys
+    expected_pocket = EMPTY_POCKET_ID
+    assert jnp.array_equal(state.player.pocket, expected_pocket), (
+        "Expected player to have {}, got {}".format(expected_pocket, state.player.pocket)
+    )
+
+    state = nx.actions.forward(state)
+
+    # check that we cannot open a door without the required key
+    state = nx.actions.open(state)
+    expected_pocket = EMPTY_POCKET_ID
+    # check that pocket is empty
+    assert jnp.array_equal(state.player.pocket, expected_pocket), (
+        "Expected player to have pocket {}, got {}".format(expected_pocket, state.player.pocket)
+    )
+    assert jnp.array_equal(state.doors.position[0], jnp.asarray((1, 3))), (
+        "Expected door position to be {}, got {}".format((1, 3), state.doors.position)
+    )
+
+    # artificially put the right key in the player's pocket
+    state = state.replace(player=state.player.replace(pocket=jnp.asarray(1)))
+
+    # check that we can open the door
+    state = nx.actions.open(state)
+    assert jnp.array_equal(state.doors.position[0], DISCARD_PILE_COORDS), (
+        "Expected door position to be {}, got {}".format(DISCARD_PILE_COORDS, state.doors.position)
+    )
+    expected_pocket = jnp.asarray(1)
+    assert jnp.array_equal(state.player.pocket, expected_pocket), (
+        "Expected player to have pocket {}, got {}".format(expected_pocket, state.player.pocket)
+    )
+
 
 if __name__ == "__main__":
     # test_rotation()
-    test_walkable()
+    # test_walkable()
+    test_pickup()
diff --git a/tests/test_environments.py b/tests/test_environments.py
index a7020cb..69a5987 100644
--- a/tests/test_environments.py
+++ b/tests/test_environments.py
@@ -73,4 +73,4 @@ def f():
 
 if __name__ == "__main__":
     test_room()
-    # test_keydoor()
+    test_keydoor()
diff --git a/tests/test_grid.py b/tests/test_grid.py
index f402be9..5b272d4 100644
--- a/tests/test_grid.py
+++ b/tests/test_grid.py
@@ -1,10 +1,9 @@
 import jax
 import jax.numpy as jnp
 import navix as nx
-from navix.grid import two_rooms
 
 
-def test_grid():
+def test_grid_from_ascii():
     ascii_map = """########
     #1.....#
     #......#
@@ -30,43 +29,37 @@ def test_grid():
 
 
 def test_idx_from_coordinates():
-    height = 2
-    width = 3
-    grid = two_rooms(height=height, width=width, key=jax.random.PRNGKey(7))[0]
+    grid = jnp.zeros((5, 7), dtype=jnp.int32)
+    grid = jnp.pad(grid, 1, mode="constant", constant_values=-1)
 
-    coordinates = jnp.mgrid[0:width, 0:height].reshape(2, -1).T
-    coordinates = jnp.asarray(coordinates, dtype=jnp.int32)
-
-    def _test(coord):
-        idx = nx.grid.idx_from_coordinates(grid, coord)
-        coord_after = nx.grid.coordinates_from_idx(grid, idx)
-        return coord_after
-
-    coordinates_after = jax.vmap(_test)(coordinates)
-
-    assert jnp.all(jnp.array_equal(coordinates, coordinates_after)), (
-        coordinates,
-        coordinates_after,
+    positions = jnp.array([[1, 1], [2, 2], [3, 3], [4, 4]])
+    indices = nx.grid.idx_from_coordinates(grid, positions)
+    positions_after = nx.grid.coordinates_from_idx(grid, indices)
+    assert jnp.all(jnp.array_equal(positions, positions_after)), (
+        positions,
+        positions_after,
     )
 
 
 def test_random_positions():
-    def f():
-        env = nx.environments.KeyDoor(height=6, width=18, max_steps=100)
-        key = jax.random.PRNGKey(7)
-        timestep = env.reset(key)
-        # without the `exclude` params in `random_positions` this
-        # specific configuration of seed and (width, height) draws player
-        # pos [2, 1] and key pos [2, 1] check that this does not happen anymore
-        return timestep.state.player.position, timestep.state.keys.position[0]
+    grid = jnp.zeros((5, 7), dtype=jnp.int32)
+    grid = jnp.pad(grid, 1, mode="constant", constant_values=-1)
+
+    key = jax.random.PRNGKey(0)
+    positions = nx.grid.random_positions(key, grid, n=1)
+    assert positions.shape == (2,), positions.shape
+
+    positions = nx.grid.random_positions(key, grid, n=4)
+    assert positions.shape == (4, 2), positions.shape
 
-    player_pos, key_pos = f()
-    assert not jnp.array_equal(player_pos, key_pos)
-    player_pos, key_pos = jax.jit(f)()
-    assert not jnp.array_equal(player_pos, key_pos)
+    exclude = jnp.asarray((1, 1))
+    positions = nx.grid.random_positions(key, grid, n=50, exclude=exclude)
+    for position in positions:
+        assert not jnp.array_equal(position, exclude), position
+        assert jnp.array_equal(grid[tuple(position)], 0), positions
 
 
 if __name__ == "__main__":
-    # test_grid()
+    test_grid_from_ascii()
+    test_idx_from_coordinates()
     test_random_positions()
-    # test_idx_from_coordinates()
diff --git a/tests/test_observations.py b/tests/test_observations.py
index fe5efd2..8c2f38f 100644
--- a/tests/test_observations.py
+++ b/tests/test_observations.py
@@ -1,14 +1,56 @@
 import jax
+import jax.numpy as jnp
+
 import navix as nx
 
 
 def test_rgb():
-    env = nx.environments.KeyDoor(
-        height=10, width=5, max_steps=100, observation_fn=nx.observations.rgb
+    height = 10
+    width = 10
+    grid = jnp.zeros((height - 2, width - 2), dtype=jnp.int32)
+    grid = jnp.pad(grid, 1, mode="constant", constant_values=-1)
+
+    state = nx.components.State(
+        key=jax.random.PRNGKey(0),
+        grid=grid,
+        player=nx.components.Player(position=jnp.asarray((1, 1))),
+        goals=nx.components.Goal(position=jnp.asarray((4, 4))[None]),
+        keys=nx.components.Key(position=jnp.asarray((2, 2))[None]),
+        doors=nx.components.Door(position=jnp.asarray((1, 5))[None]),
+        cache=nx.graphics.RenderingCache.init(grid),
     )
-    key = jax.random.PRNGKey(4)
-    state = jax.jit(env.reset)(key)
-    print(state.observation)
+    void = jnp.zeros((nx.graphics.TILE_SIZE, nx.graphics.TILE_SIZE, 3), dtype=jnp.uint8)
+    tiles_registry={
+        "player": void,
+        "goal": void,
+        "key": void,
+        "door": void,
+    }
+
+    obs = nx.observations.rgb(state, tiles_registry=tiles_registry)
+    expected_obs_shape = (height * nx.graphics.TILE_SIZE, width * nx.graphics.TILE_SIZE, 3)
+    assert obs.shape == expected_obs_shape, (
+        f"Expected observation {expected_obs_shape}, got {obs.shape} instead"
+    )
+
+    def get_tile(position):
+        x = position[0] * nx.graphics.TILE_SIZE
+        y = position[1] * nx.graphics.TILE_SIZE
+        return obs[x:x + nx.graphics.TILE_SIZE, y:y + nx.graphics.TILE_SIZE, :]
+
+    player_tile = get_tile(state.player.position)
+    assert jnp.array_equal(player_tile, tiles_registry["player"]), player_tile
+
+    goal_tile = get_tile(state.goals.position[0])
+    assert jnp.array_equal(goal_tile, tiles_registry["goal"]), goal_tile
+
+    key_tile = get_tile(state.keys.position[0])
+    assert jnp.array_equal(key_tile, tiles_registry["key"]), key_tile
+
+    door_tile = get_tile(state.doors.position[0])
+    assert jnp.array_equal(door_tile, tiles_registry["door"]), door_tile
+
+    return
 
 
 if __name__ == "__main__":