dfa-gym 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

dfa_gym/token_env.py

@@ -0,0 +1,571 @@
+import jax
+import chex
+import numpy as np
+import jax.numpy as jnp
+from flax import struct
+from enum import IntEnum
+from functools import partial
+from typing import Tuple, Dict
+from dfa_gym import spaces
+from dfa_gym.env import MultiAgentEnv, State
+
+
+class Action(IntEnum):
+    DOWN  = 0
+    RIGHT = 1
+    UP    = 2
+    LEFT  = 3
+    NOOP  = 4
+
+ACTION_MAP = jnp.array([
+    [ 1,  0], # DOWN
+    [ 0,  1], # RIGHT
+    [-1,  0], # UP
+    [ 0, -1], # LEFT
+    [ 0,  0], # NOOP
+])
+
+@struct.dataclass
+class TokenEnvState(State):
+    agent_positions: jax.Array
+    token_positions: jax.Array
+    wall_positions: jax.Array
+    is_wall_disabled: jax.Array
+    button_positions: jax.Array
+    is_alive: jax.Array
+    time: int
+
+class TokenEnv(MultiAgentEnv):
+
+    def __init__(
+        self,
+        n_agents: int = 3,
+        n_tokens: int = 10,
+        n_token_repeat: int = 2,
+        grid_shape: Tuple[int, int] = (7, 7),
+        fixed_map_seed: int | None = None,
+        max_steps_in_episode: int = 100,
+        collision_reward: int | None = None,
+        black_death: bool = True,
+        layout: str | None = None
+    ) -> None:
+        super().__init__(num_agents=n_agents)
+        assert (grid_shape[0] * grid_shape[1]) >= (n_agents + n_tokens * n_token_repeat)
+        self.n_agents = n_agents
+        self.n_tokens = n_tokens
+        self.n_token_repeat = n_token_repeat
+        self.grid_shape = grid_shape
+        self.grid_shape_arr = jnp.array(self.grid_shape)
+        self.fixed_map_seed = fixed_map_seed
+        self.max_steps_in_episode = max_steps_in_episode
+        self.collision_reward = collision_reward
+        self.black_death = black_death
+        self.n_buttons = 0
+
+        self.agents = [f"agent_{i}" for i in range(self.n_agents)]
+
+        self.init_state = None
+        if layout is not None:
+            self.init_state = self.parse(layout)
+        self.num_agents = self.n_agents
+
+        channel_dim = 1
+        if self.init_state is not None: channel_dim += 2
+        if self.n_tokens > 0: channel_dim += self.n_tokens
+        if self.n_agents > 1: channel_dim += self.n_agents - 1
+        if self.n_buttons > 0: channel_dim += 3 * self.n_buttons
+        self.obs_shape = (channel_dim, *self.grid_shape)
+
+        self.action_spaces = {
+            agent: spaces.Discrete(len(ACTION_MAP))
+            for agent in self.agents
+        }
+        self.observation_spaces = {
+            agent: spaces.Box(low=0, high=1, shape=self.obs_shape, dtype=jnp.uint8)
+            for agent in self.agents
+        }
+
+    @partial(jax.jit, static_argnums=(0,))
+    def reset(
+        self,
+        key: chex.PRNGKey
+    ) -> Tuple[Dict[str, chex.Array], TokenEnvState]:
+        state = self.init_state
+        if state is None: state = self.sample_init_state(key)
+        obs = self.get_obs(state=state)
+        return obs, state
+
+    @partial(jax.jit, static_argnums=(0,))
+    def step_env(
+        self,
+        key: chex.PRNGKey,
+        state: TokenEnvState,
+        actions: Dict[str, chex.Array]
+    ) -> Tuple[Dict[str, chex.Array], TokenEnvState, Dict[str, float], Dict[str, bool], Dict]:
+
+        _actions = jnp.array([actions[agent] for agent in self.agents])
+
+        # Move agents
+        def move_agent(pos, a):
+            return (pos + ACTION_MAP[a]) % self.grid_shape_arr
+        new_agent_pos = jax.vmap(move_agent, in_axes=(0, 0))(state.agent_positions, _actions)
+        new_agent_pos = jnp.where(state.is_alive[:, None], new_agent_pos, state.agent_positions)
+
+        if self.init_state is not None:
+            # Handle wall collisions
+            def compute_wall_collisions(pos, wall_positions, is_wall_disabled):
+                return jnp.any(
+                    jnp.logical_and(
+                        jnp.logical_not(is_wall_disabled),
+                        jnp.all(
+                            pos[None, :] == wall_positions # [N, 2]
+                        , axis=-1), # [N,]
+                    )
+                , axis=-1) # [1,]
+            wall_collisions = jax.vmap(compute_wall_collisions, in_axes=(0, None, None))(new_agent_pos, state.wall_positions, state.is_wall_disabled)
+            new_agent_pos = jnp.where(wall_collisions[:, None], state.agent_positions, new_agent_pos)
+
+        # Handle collisions
+        # TODO: When collision_reward is not None, there might be unintended behavior.
+        # +-------+-------+-------+-------+-------+-------+-------+
+        # | 0     | #     | .     | #     | .     | #     | 2     |
+        # +-------+-------+-------+-------+-------+-------+-------+
+        # | .     | #     | .     | #     | 1     | #     | .     |
+        # +-------+-------+-------+-------+-------+-------+-------+
+        # | 4     | #     | 9     | 5     | 3     | #     | 7     |
+        # +-------+-------+-------+-------+-------+-------+-------+
+        # | .     | .     | 8     | #     | .     | #     | 7     |
+        # +-------+-------+-------+-------+-------+-------+-------+
+        # | 6     | #     | 2     | #     | 9     | #     | 3     |
+        # +-------+-------+-------+-------+-------+-------+-------+
+        # | 8     | #     | 0     | #     | .     | #     | A_1,5 |
+        # +-------+-------+-------+-------+-------+-------+-------+
+        # | .     | #     | 1     | #     | 4     | A_2   | A_0,6 |
+        # +-------+-------+-------+-------+-------+-------+-------+
+        # Action for agent_0
+        # 3
+        # Action for agent_1
+        # 0
+        # Action for agent_2
+        # 1
+        # Gives
+        # {'agent_0': Array(-100., dtype=float32), 'agent_1': Array(-100., dtype=float32), 'agent_2': Array(-100., dtype=float32)}
+        # {'__all__': Array(True, dtype=bool), 'agent_0': Array(True, dtype=bool), 'agent_1': Array(True, dtype=bool), 'agent_2': Array(True, dtype=bool)}
+        def compute_collisions(mask):
+            positions = jnp.where(mask[:, None], state.agent_positions, new_agent_pos)
+
+            collision_grid = jnp.zeros(self.grid_shape)
+            collision_grid, _ = jax.lax.scan(
+                lambda grid, pos: (grid.at[pos[0], pos[1]].add(1), None),
+                collision_grid,
+                positions,
+            )
+
+            collision_mask = collision_grid > 1
+
+            collisions = jax.vmap(lambda p: collision_mask[p[0], p[1]])(positions)
+            return jnp.logical_and(state.is_alive, collisions)
+
+        collisions = jax.lax.while_loop(
+            lambda mask: jnp.any(compute_collisions(mask)),
+            lambda mask: jnp.logical_or(mask, compute_collisions(mask)),
+            jnp.zeros((self.n_agents,), dtype=bool)
+        )
+
+        if self.collision_reward is None:
+            new_agent_pos = jnp.where(collisions[:, None], state.agent_positions, new_agent_pos)
+            collisions = jnp.full(collisions.shape, False)
+
+        # Handle swaps
+        def compute_swaps(original_positions, new_positions):
+            original_pos_expanded = jnp.expand_dims(original_positions, axis=0)
+            new_pos_expanded = jnp.expand_dims(new_positions, axis=1)
+
+            swap_mask = (original_pos_expanded == new_pos_expanded).all(axis=-1)
+            swap_mask = jnp.fill_diagonal(swap_mask, False, inplace=False)
+
+            swap_pairs = jnp.logical_and(swap_mask, swap_mask.T)
+
+            swaps = jnp.any(swap_pairs, axis=0)
+            return swaps
+
+        swaps = compute_swaps(state.agent_positions, new_agent_pos)
+        new_agent_pos = jnp.where(swaps[:, None], state.agent_positions, new_agent_pos)
+
+        _rewards = jnp.zeros((self.n_agents,), dtype=jnp.float32)
+        if self.collision_reward is not None:
+            _rewards = jnp.where(jnp.logical_and(state.is_alive, collisions), self.collision_reward, _rewards)
+        rewards = {agent: _rewards[i] for i, agent in enumerate(self.agents)}
+
+        is_wall_disabled = jnp.empty((0, 2), dtype=bool)
+        if self.init_state is not None:
+            is_wall_disabled = self.compute_disabled_walls(new_agent_pos, state.wall_positions, state.button_positions)
+
+        new_state = TokenEnvState(
+            agent_positions=new_agent_pos,
+            token_positions=state.token_positions,
+            wall_positions=state.wall_positions,
+            is_wall_disabled=is_wall_disabled,
+            button_positions=state.button_positions,
+            is_alive=jnp.logical_and(state.is_alive, jnp.logical_not(collisions)),
+            time=state.time + 1
+        )
+
+        _dones = jnp.logical_or(collisions, new_state.time >= self.max_steps_in_episode)
+        dones = {a: _dones[i] for i, a in enumerate(self.agents)}
+        dones.update({"__all__": jnp.all(_dones)})
+
+        obs = self.get_obs(state=new_state)
+        info = {}
+
+        return obs, new_state, rewards, dones, info
+
+    @partial(jax.jit, static_argnums=(0,))
+    def get_obs(
+        self,
+        state: TokenEnvState
+    ) -> Dict[str, chex.Array]:
+
+        def obs_for_agent(i):
+            base = jnp.zeros(self.obs_shape, dtype=jnp.uint8)
+            # ref = self.grid_shape_arr // 2
+            ref = jnp.array([0, 0])
+            offset = ref - state.agent_positions[i]
+            idx_offset = 0
+            b = base
+
+            def place_agent(val):
+                rel = (state.agent_positions[i] + offset) % self.grid_shape_arr
+                return val.at[idx_offset, rel[0], rel[1]].set(1) # Is agent?
+            b = place_agent(b)
+            idx_offset += 1
+
+            if self.init_state is not None:
+                def place_wall(val):
+                    rel = (state.wall_positions + offset) % self.grid_shape_arr
+                    return val.at[
+                        idx_offset, rel[:, 0], rel[:, 1]
+                    ].set(1).at[ # Is wall?
+                        idx_offset + 1, rel[:, 0], rel[:, 1]
+                    ].set(jnp.logical_not(state.is_wall_disabled).astype(jnp.uint8)) # Is wall blocking?
+                b = place_wall(b)
+                idx_offset += 2
+
+            if self.n_tokens > 0:
+                def place_token(token_idx, val):
+                    rel = (state.token_positions[token_idx] + offset) % self.grid_shape_arr
+                    return val.at[idx_offset + token_idx, rel[:, 0], rel[:, 1]].set(1) # Is token?
+                b = jax.lax.fori_loop(0, self.n_tokens, place_token, b)
+                idx_offset += self.n_tokens
+
+            if self.n_agents > 1:
+                def place_other(other_idx, val):
+                    rel = (state.agent_positions[other_idx + (other_idx >= i)] + offset) % self.grid_shape_arr
+                    return val.at[idx_offset + other_idx, rel[0], rel[1]].set(1) # Is other agent?
+                b = jax.lax.fori_loop(0, self.n_agents - 1, place_other, b)
+                idx_offset += self.n_agents - 1
+
+            if self.n_buttons > 0:
+                def place_button(button_idx, val):
+                    is_door = jnp.any(
+                        jnp.all(
+                            state.button_positions[button_idx][:, None, :] == state.wall_positions[None, :, :]
+                        , axis=-1)
+                    , axis=-1) # Buttons are considered doors if they are in a wall.
+                    rel = (state.button_positions[button_idx] + offset) % self.grid_shape_arr
+                    return val.at[
+                        idx_offset + 3 * button_idx, rel[:, 0], rel[:, 1]
+                    ].set(1).at[ # Is button-door pair?
+                        idx_offset + 3 * button_idx + 1, rel[:, 0], rel[:, 1]
+                    ].set(jnp.logical_not(is_door).astype(jnp.uint8)).at[ # Is button?
+                        idx_offset + 3 * button_idx + 2, rel[:, 0], rel[:, 1]
+                    ].set(is_door.astype(jnp.uint8)) # Is door?
+
+                b = jax.lax.fori_loop(0, self.n_buttons, place_button, b)
+
+            return jnp.where(jnp.logical_or(jnp.logical_not(self.black_death), state.is_alive[i]), b, base)
+
+        obs = jax.vmap(obs_for_agent)(jnp.arange(self.n_agents))
+        return {agent: obs[i] for i, agent in enumerate(self.agents)}
+
+    @partial(jax.jit, static_argnums=(0,))
+    def label_f(self, state: TokenEnvState) -> Dict[str, int]:
+
+        diffs = state.agent_positions[:, None, None, :] - state.token_positions[None, :, :, :]
+        matches = jnp.all(diffs == 0, axis=-1)
+        matches_any = jnp.any(matches, axis=-1)
+
+        has_match = jnp.any(matches_any, axis=1)
+        token_idx = jnp.argmax(matches_any, axis=1)
+
+        agent_token_matches = jnp.where(jnp.logical_and(has_match, state.is_alive), token_idx, -1)
+
+        return {self.agents[agent_idx]: token_idx for agent_idx, token_idx in enumerate(agent_token_matches)}
+
+    @partial(jax.jit, static_argnums=(0,))
+    def sample_init_state(
+        self,
+        key: chex.PRNGKey
+    ) -> Tuple[Dict[str, chex.Array], TokenEnvState]:
+        if self.fixed_map_seed is not None:
+            key = jax.random.PRNGKey(self.fixed_map_seed)
+
+        grid_points = jnp.stack(jnp.meshgrid(jnp.arange(self.grid_shape[0]), jnp.arange(self.grid_shape[1])), -1)
+        grid_flat = grid_points.reshape(-1, 2)
+
+        key, subkey = jax.random.split(key)
+        perm = jax.random.permutation(subkey, grid_flat.shape[0])
+
+        agent_positions = grid_flat[perm][:self.n_agents]
+        token_positions = grid_flat[perm][self.n_agents: self.n_agents + self.n_tokens * self.n_token_repeat].reshape(self.n_tokens, self.n_token_repeat, 2)
+
+        return TokenEnvState(
+            agent_positions=agent_positions,
+            token_positions=token_positions,
+            wall_positions=jnp.empty((0, 2), dtype=jnp.int32),
+            is_wall_disabled=jnp.empty((0, 2), dtype=bool),
+            button_positions=jnp.empty((0, 2), dtype=jnp.int32),
+            is_alive=jnp.ones((self.n_agents,), dtype=bool),
+            time=0
+        )
+
+    @partial(jax.jit, static_argnums=(0,))
+    def compute_disabled_walls(self, agent_positions, wall_positions, button_positions):
+        def _compute_on_buttons(button_pos, agent_positions):
+            return jnp.any(
+                jnp.any(
+                    jnp.all(
+                        button_pos[:, None, :] == agent_positions[None, :, :] # [M, N, 2]
+                    , axis=-1) # [M, N,]
+                , axis=-1) # [M,]
+            , axis=-1) # [1,]
+        on_buttons = jax.vmap(_compute_on_buttons, in_axes=(0, None))(button_positions, agent_positions)
+        def _compute_disabled_walls(wall_pos, on_buttons, button_positions):
+            # Compare each wall_pos to each button coordinate
+            eq = jnp.all(button_positions == wall_pos, axis=-1)  # (n_buttons, n_button_repeat)
+            # A wall is disabled if *any* matching button is pressed
+            return jnp.any(jnp.logical_and(on_buttons, jnp.any(eq, axis=-1)))
+        return jax.vmap(_compute_disabled_walls, in_axes=(0, None, None))(wall_positions, on_buttons, button_positions)
+
+    def parse(self, layout: str) -> TokenEnvState:
+        # Example layout:
+        # [ 8 ][   ][   ][   ][   ][   ][   ][ # ][ 0 ][   ][   ][ 1 ]
+        # [   ][   ][   ][   ][   ][   ][   ][ # ][   ][   ][   ][   ]
+        # [   ][   ][ b ][   ][   ][   ][   ][ # ][   ][   ][   ][   ]
+        # [   ][   ][   ][   ][   ][   ][   ][ # ][ 3 ][   ][   ][ 2 ]
+        # [   ][   ][   ][   ][   ][   ][   ][ # ][ # ][ # ][#,a][ # ]
+        # [ A ][   ][   ][   ][   ][   ][   ][   ][   ][   ][   ][   ]
+        # [ B ][   ][   ][   ][   ][   ][   ][   ][   ][   ][   ][   ]
+        # [   ][   ][   ][   ][   ][   ][   ][ # ][ # ][ # ][#,b][ # ]
+        # [   ][   ][   ][   ][   ][   ][   ][ # ][ 4 ][   ][   ][ 5 ]
+        # [   ][   ][ a ][   ][   ][   ][   ][ # ][   ][   ][   ][   ]
+        # [   ][   ][   ][   ][   ][   ][   ][ # ][   ][   ][   ][   ]
+        # [ 9 ][   ][   ][   ][   ][   ][   ][ # ][ 7 ][   ][   ][ 6 ]
+        # where each [] indicates a cell, uppercase letters indicate agents,
+        # e.g., A and B, and lower case letters indicate "sync" points which
+        # are more like doors and buttons that open the doors, eg [#,a] is a
+        # door that is open if there is an agent on the cell [ a ] and closed;
+        # otherwise, and cells with # indicate walls.
+
+        # --- parse into a 2D list of cell contents (strings inside brackets) ---
+        lines = [ln.strip() for ln in layout.strip().splitlines() if ln.strip()]
+        rows: list[list[str]] = []
+
+        for ln in lines:
+            cells = []
+            i = 0
+            while i < len(ln):
+                if ln[i] == "[":
+                    j = ln.find("]", i + 1)
+                    if j == -1:
+                        raise ValueError("Malformed layout: missing closing ']' in a row.")
+                    cells.append(ln[i + 1:j].strip())
+                    i = j + 1
+                else:
+                    i += 1
+            if cells:
+                rows.append(cells)
+
+        if not rows:
+            raise ValueError("Parsed layout is empty.")
+
+        H = len(rows)
+        W = len(rows[0])
+        if any(len(r) != W for r in rows):
+            raise ValueError("All rows in the layout must have the same number of cells.")
+
+        # --- collect info ---
+        wall_positions: list[tuple[int, int]] = []
+        agent_positions: dict[int, tuple[int, int]] = {}
+        token_positions: dict[int, list[tuple[int, int]]] = {}
+        button_positions: dict[int, list[tuple[int, int]]] = {}
+        max_token_id = -1
+        max_button_id = -1
+
+        def _is_wall(c: str) -> bool:
+            return "#" == c
+
+        def _is_agent(c: str) -> bool:
+            is_agent_in_cell = [
+                a == c
+                for a in [chr(ord("A") + i) for i in range(ord("Z") - ord("A") + 1)]
+            ]
+            return sum(is_agent_in_cell) > 0
+
+        def _get_agent_idx(c: str) -> bool:
+            is_agent_in_cell = [
+                a == c
+                for a in [chr(ord("A") + i) for i in range(ord("Z") - ord("A") + 1)]
+            ]
+            return np.argmax(is_agent_in_cell)
+
+        def _is_token(c: str) -> bool:
+            try:
+                int(c)
+                return True
+            except ValueError:
+                return False
+
+        def _is_button(c: str) -> bool:
+            is_button_in_cell = [
+                a == c
+                for a in [chr(ord("a") + i) for i in range(ord("z") - ord("a") + 1)]
+            ]
+            return sum(is_button_in_cell) > 0
+
+        def _get_button_idx(c: str) -> bool:
+            is_button_in_cell = [
+                a == c
+                for a in [chr(ord("a") + i) for i in range(ord("z") - ord("a") + 1)]
+            ]
+            return np.argmax(is_button_in_cell)
+
+        for r in range(H):
+            for c in range(W):
+                cell = rows[r][c]
+                has_wall = False
+                has_agent = False
+                has_token = False
+                has_button = False
+                for content in cell.split(","):
+
+                    if _is_wall(content):
+                        if has_wall:
+                            raise ValueError(f"One wall per cell.")
+                        has_wall = True
+                        wall_positions.append((r, c))
+
+                    if _is_agent(content):
+                        if has_agent:
+                            raise ValueError(f"One agent per cell.")
+                        has_agent = True
+                        idx = _get_agent_idx(content)
+                        if idx in agent_positions:
+                            raise ValueError(f"Duplicate placement for agent '{content}'.")
+                        agent_positions[idx] = (r, c)
+
+                    if _is_token(content):
+                        if has_token:
+                            raise ValueError(f"One token per cell.")
+                        has_token = True
+                        tok_id = int(content)
+                        max_token_id = max(max_token_id, tok_id)
+                        token_positions.setdefault(tok_id, []).append((r, c))
+
+                    if _is_button(content):
+                        if has_button:
+                            raise ValueError(f"One button per cell.")
+                        has_button = True
+                        idx = _get_button_idx(content)
+                        max_button_id = max(max_button_id, idx)
+                        button_positions.setdefault(idx, []).append((r, c))
+
+                assert not (has_wall and has_agent)
+                assert not (has_wall and has_token)
+                assert not (has_token and has_button)
+
+        # --- override environment settings ---
+        self.grid_shape = (H, W)
+        self.grid_shape_arr = jnp.array(self.grid_shape)
+
+        self.n_agents = len(agent_positions)
+        self.agents = [f"agent_{i}" for i in range(self.n_agents)]
+
+        self.n_tokens = max_token_id + 1 if max_token_id >= 0 else 0
+        self.n_token_repeat = max((len(v) for v in token_positions.values()), default=0)
+        token_positions_np = np.full((self.n_tokens, self.n_token_repeat, 2), -1, dtype=np.int32)
+        for tid in range(self.n_tokens):
+            coords = token_positions.get(tid, [])
+            for k, (r, c) in enumerate(coords[: self.n_token_repeat]):
+                token_positions_np[tid, k] = (r, c)
+
+        agent_positions_np = np.full((self.n_agents, 2), -1, dtype=np.int32)
+        for idx, pos in agent_positions.items():
+            agent_positions_np[idx] = pos
+
+        wall_positions_np = np.array(wall_positions, dtype=np.int32) if wall_positions else np.empty((0, 2), dtype=np.int32)
+
+        self.n_buttons = max_button_id + 1 if max_button_id >= 0 else 0
+        self.n_button_repeat = max((len(v) for v in button_positions.values()), default=0)
+        button_positions_np = np.full((self.n_buttons, self.n_button_repeat, 2), -1, dtype=np.int32)
+        for bid in range(self.n_buttons):
+            coords = button_positions.get(bid, [])
+            for k, (r, c) in enumerate(coords[: self.n_button_repeat]):
+                button_positions_np[bid, k] = (r, c)
+
+        agent_positions_jnp = jnp.array(agent_positions_np)
+        wall_positions_jnp = jnp.array(wall_positions_np)
+        button_positions_jnp = jnp.array(button_positions_np)
+
+        # --- return state ---
+        return TokenEnvState(
+            agent_positions=agent_positions_jnp,
+            token_positions=jnp.array(token_positions_np),
+            wall_positions=wall_positions_jnp,
+            is_wall_disabled=self.compute_disabled_walls(agent_positions_jnp, wall_positions_jnp, button_positions_jnp),
+            button_positions=button_positions_jnp,
+            is_alive=jnp.ones((self.n_agents,), dtype=bool),
+            time=0,
+        )
+
+    def render(self, state: TokenEnvState):
+        empty_cell = "."
+        wall_cell = "#"
+        grid = np.full(self.grid_shape, empty_cell, dtype=object)
+
+        for disabled, pos in zip(state.is_wall_disabled, state.wall_positions):
+            if not disabled:
+                grid[pos[0], pos[1]] = f"{wall_cell}"
+
+        for token, positions in enumerate(state.token_positions):
+            for pos in positions:
+                grid[pos[0], pos[1]] = f"{token}"
+
+        for button, positions in enumerate(state.button_positions):
+            for pos in positions:
+                current = grid[pos[0], pos[1]]
+                if current == empty_cell:
+                    grid[pos[0], pos[1]] = f"b_{button}"
+                else:
+                    grid[pos[0], pos[1]] = f"{current},b_{button}"
+
+        for agent in range(self.n_agents):
+            pos = state.agent_positions[agent]
+            current = grid[pos[0], pos[1]]
+            if current == empty_cell:
+                grid[pos[0], pos[1]] = f"A_{agent}"
+            else:
+                grid[pos[0], pos[1]] = f"A_{agent},{current}"
+
+        max_width = max(len(str(cell)) for row in grid for cell in row)
+
+        out = ""
+        h_line = "+" + "+".join(["-" * (max_width + 2) for _ in range(self.grid_shape[1])]) + "+"
+        out += h_line + "\n"
+        for row in grid:
+            row_str = "| " + " | ".join(f"{str(cell):<{max_width}}" for cell in row) + " |"
+            out += row_str + "\n"
+            out += h_line + "\n"
+
+        print(out)
+