mapox 0.1.0__py3-none-any.whl

mapox/__init__.py

@@ -0,0 +1,30 @@
+"""MAPOX: Multi-Agent Partially Observable gridworlds in JAX"""
+
+__version__ = "0.1.0"
+
+from mapox.timestep import TimeStep
+from mapox.environment import Environment
+from mapox.specs import ActionSpec, ObservationSpec
+from mapox.wrappers.multitask import MultiTaskWrapper
+from mapox.wrappers.vector import VectorWrapper
+from mapox.client import GridworldClient
+from mapox.utils.encode_one_hot import concat_one_hot
+from mapox.config import create_env, create_client, ScoutsConfig, TravelingSalesmanConfig, KingHillConfig, MultiTaskConfig, EnvironmentConfig
+
+__all__ = [
+    "TimeStep",
+    "ActionSpec",
+    "ObservationSpec",
+    "Environment",
+    "MultiTaskWrapper",
+    "VectorWrapper",
+    "GridworldClient",
+    "concat_one_hot",
+    "create_env",
+    "create_client",
+    "ScoutsConfig",
+    "TravelingSalesmanConfig",
+    "KingHillConfig",
+    "MultiTaskConfig",
+    "EnvironmentConfig"
+]

mapox/client.py

@@ -0,0 +1,53 @@
+import pygame
+from abc import ABC, abstractmethod
+
+from mapox.renderer import GridworldRenderer, GridRenderState
+from mapox.timestep import TimeStep
+
+
+class EnvironmentClient[State](ABC):
+    @abstractmethod
+    def render(self, state: State, timestep: TimeStep): ...
+
+    @abstractmethod
+    def save_video(self): ...
+
+
+class GridworldClient:
+    """EnvironmentClient that renders via GridworldRenderer using per-env adapters."""
+
+    def __init__(
+        self, env, screen_width: int = 960, screen_height: int = 960, fps: int = 10
+    ):
+        assert hasattr(env, "get_render_state"), (
+            "Env must implement get_render_state(state)"
+        )
+        assert hasattr(env, "get_render_settings"), (
+            "Env must implement get_render_settings()"
+        )
+        self.env = env
+        self.renderer = GridworldRenderer(
+            screen_width=screen_width, screen_height=screen_height, fps=fps
+        )
+        self.renderer.set_env(env.get_render_settings())
+
+    def render(self, state, timestep):
+        rs: GridRenderState = self.env.get_render_state(state)
+        self.renderer.render(rs)
+
+    def render_pov(self, state, timestep):
+        """Render only the focused agent's point-of-view, filling the screen."""
+        rs: GridRenderState = self.env.get_render_state(state)
+        self.renderer.render_agent_view(rs)
+
+    def handle_event(self, event: pygame.event.Event) -> bool:
+        return self.renderer.handle_event(event)
+
+    def record_frame(self):
+        self.renderer.record_frame()
+
+    def save_video(self, file_name: str):
+        self.renderer.save_video(file_name)
+
+    def focus_agent(self, agent_id: int | None):
+        self.renderer.focus_agent(agent_id)

mapox/config.py

@@ -0,0 +1,92 @@
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+from mapox.envs.king_hill import KingHillConfig, KingHillEnv
+from mapox.envs.grid_return import ReturnDiggingConfig, ReturnDiggingEnv
+from mapox.envs.traveling_salesman import (
+    TravelingSalesmanConfig,
+    TravelingSalesmanEnv,
+)
+from mapox.envs.scouts import ScoutsConfig, ScoutsEnv
+from mapox.client import GridworldClient
+
+from mapox.client import EnvironmentClient
+from mapox.environment import Environment
+from mapox.wrappers.task_id_wrapper import TaskIdWrapper
+from mapox.wrappers.multitask import MultiTaskWrapper
+from mapox.wrappers.vector import VectorWrapper
+
+type EnvironmentConfig = (
+    ReturnDiggingConfig
+    | TravelingSalesmanConfig
+    | ScoutsConfig
+    | KingHillConfig
+)
+
+
+class MultiTaskEnvConfig(BaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True)
+    num: int = 1
+    name: str
+    env: EnvironmentConfig = Field(discriminator="env_type")
+
+
+class MultiTaskConfig(BaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True)
+    env_type: Literal["multi"] = "multi"
+    envs: tuple[MultiTaskEnvConfig, ...]
+
+    @field_validator("envs", mode="before")
+    @classmethod
+    def coerce_envs(cls, v):
+        # JSON gives list; accept list and turn into tuple
+        return tuple(v) if isinstance(v, list) else v
+
+
+def create_env(
+    env_config: EnvironmentConfig | MultiTaskConfig,
+    length: int,
+    vec_count: int = 1,
+    env_name: str | None = None,
+) -> tuple[Environment, int]:
+    num_tasks = 1
+    if env_config.env_type == "multi" and env_name is not None:
+        num_tasks = len(env_config.envs)
+        for task_id, env_def in enumerate(env_config.envs):
+            if env_def.name == env_name:
+                return TaskIdWrapper(
+                    create_env(env_def.env, length, vec_count=vec_count)[0], task_id
+                ), num_tasks
+        raise ValueError("Could not find environment matching env_name")
+
+    match env_config.env_type:
+        case "multi":
+            out_envs = []
+            out_env_names = []
+            num_tasks = len(env_config.envs)
+            for env_def in env_config.envs:
+                out_envs.append(create_env(env_def.env, length, env_def.num)[0])
+                out_env_names.append(env_def.name)
+
+            env = MultiTaskWrapper(tuple(out_envs), tuple(out_env_names))
+        case "return_digging":
+            env = ReturnDiggingEnv(env_config, length)
+        case "scouts":
+            env = ScoutsEnv(env_config, length)
+        case "traveling_salesman":
+            env = TravelingSalesmanEnv(env_config, length)
+        case "king_hill":
+            env = KingHillEnv(env_config, length)
+        case _:
+            raise ValueError(f"Unknown environment type: {env_config.env_type}")
+
+    if vec_count > 1:
+        env = VectorWrapper(env, vec_count)
+
+    return env, num_tasks
+
+
+def create_client[State](env: Environment[State]) -> EnvironmentClient[State]:
+    # use the grid client for all environments for now
+    return GridworldClient(env)

mapox/environment.py

@@ -0,0 +1,52 @@
+from abc import ABC, abstractmethod
+from functools import cached_property
+from typing import Any
+
+import jax
+from jax import Array
+
+from mapox.timestep import TimeStep
+from mapox.specs import ObservationSpec, ActionSpec
+import enum
+
+
+class StepType(enum.IntEnum):
+    FIRST = 0
+    MID = 1
+    LAST = 2
+
+
+class Environment[State](ABC):
+    @abstractmethod
+    def reset(self, rng_key: Array) -> tuple[State, TimeStep]: ...
+
+    @abstractmethod
+    def step(
+        self, state: State, action: Array, rng_key: Array
+    ) -> tuple[State, TimeStep]: ...
+
+    @abstractmethod
+    def create_placeholder_logs(self) -> dict[str, Any]: ...
+
+    @abstractmethod
+    def create_logs(self, state) -> dict[str, Any]: ...
+
+    @cached_property
+    @abstractmethod
+    def observation_spec(self) -> ObservationSpec: ...
+
+    @cached_property
+    @abstractmethod
+    def action_spec(self) -> ActionSpec: ...
+
+    @property
+    @abstractmethod
+    def num_agents(self) -> int: ...
+
+    @property
+    @abstractmethod
+    def is_jittable(self) -> bool: ...
+
+    @property
+    def teams(self) -> jax.Array | None:
+        return None

mapox/envs/constance.py

@@ -0,0 +1,75 @@
+from jax import numpy as jnp
+
+from mapox.specs import ObservationSpec
+
+NUM_TYPES = 15
+
+# Unified tile ids across gridworld environments
+TILE_EMPTY = 0  # empty space
+TILE_WALL = 1  # permanent wall
+TILE_DESTRUCTIBLE_WALL = 2  # destructible
+TILE_FLAG = 3  # typical goal tile
+TILE_FLAG_UNLOCKED = 4  # used for the scouting environment where the the flag gets made available for taking
+
+# TILE_FLAG_BLUE_TEAM = 5
+# TILE_FLAG_RED_TEAM = 6
+
+# agents are observed like tiles
+AGENT_GENERIC = 5  # typical agent
+AGENT_SCOUT = 6  # scout agent (scout env)
+AGENT_HARVESTER = 7  # harvester agent (scout env)
+
+AGENT_KNIGHT = 8
+AGENT_ARCHER = 9
+
+TILE_DECOR_1 = 10
+TILE_DECOR_2 = 11
+TILE_DECOR_3 = 12
+TILE_DECOR_4 = 13
+
+TILE_ARROW = 14
+
+
+# Actions
+NUM_ACTIONS = 7
+
+MOVE_UP = 0
+MOVE_RIGHT = 1
+MOVE_DOWN = 2
+MOVE_LEFT = 3
+STAY = 4
+PRIMARY_ACTION = 5
+DIG_ACTION = 6
+
+
+DIRECTIONS = jnp.array([[0, 1], [1, 0], [0, -1], [-1, 0]], dtype=jnp.int32)
+
+
+def make_obs_spec(width: int, height: int) -> ObservationSpec:
+    # CHANNELS:
+    # TILE
+    # DIRECTION
+    # TEAM ID
+    # HEALTH
+    obs_spec = ObservationSpec(
+        dtype=jnp.int8,
+        shape=(width, height, 4),
+        max_value=(
+            NUM_TYPES,
+            5,  # none, up, right, down, left,
+            3,  # none, red, blue
+            3,  # 0, 1, 2
+        ),
+    )
+
+    return obs_spec
+
+
+def make_action_mask(actions: list[int], num_agents: int):
+    mask = [False] * NUM_ACTIONS
+
+    for action in actions:
+        mask[action] = True
+
+    mask_array = jnp.array(mask, jnp.bool)
+    return jnp.repeat(mask_array[None, :], num_agents, axis=0)

mapox/envs/grid_return.py

@@ -0,0 +1,389 @@
+from functools import cached_property, partial
+from typing import NamedTuple, Literal
+
+import jax
+from jax import numpy as jnp
+import numpy as np
+from pydantic import BaseModel, ConfigDict
+
+from mapox.map_generator import (
+    fractal_noise,
+    generate_decor_tiles,
+    choose_positions,
+)
+from mapox.environment import Environment
+from mapox.specs import DiscreteActionSpec, ObservationSpec
+from mapox.timestep import TimeStep
+from mapox.envs.renderer import GridRenderSettings, GridRenderState
+import mapox.envs.constance as GW
+
+
+class ReturnDiggingConfig(BaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True)
+    env_type: Literal["return_digging"] = "return_digging"
+
+    num_agents: int = 1
+    num_flags: int = 1
+
+    width: int = 40
+    height: int = 40
+    view_width: int = 11
+    view_height: int = 11
+
+    mapgen_threshold: float = 0.3
+    digging_timeout: int = 5
+    treasure_reward: float = 1.0
+
+    eval_map: bool = False
+
+
+class ReturnDiggingState(NamedTuple):
+    agents_pos: jax.Array
+    agents_timeout: jax.Array
+    found_reward: jax.Array
+
+    time: jax.Array
+
+    map: jax.Array
+    spawn_pos: jax.Array
+    spawn_count: jax.Array
+
+    rewards: jax.Array
+
+
+class ReturnDiggingEnv(Environment[ReturnDiggingState]):
+    def __init__(self, config: ReturnDiggingConfig, length: int) -> None:
+        super().__init__()
+
+        self._config = config
+        self._length = length
+        self._num_agents = config.num_agents
+        self.num_flags = config.num_flags
+
+        self.unpadded_width = config.width
+        self.unpadded_height = config.height
+
+        self.view_width = config.view_width
+        self.view_height = config.view_height
+        self.pad_width = self.view_width // 2
+        self.pad_height = self.view_height // 2
+
+        self.width = self.unpadded_width + self.pad_width
+        self.height = self.unpadded_height + self.pad_height
+
+        self.mapgen_threshold = config.mapgen_threshold
+        self.digging_timeout = config.digging_timeout
+        self.treasure_reward = config.treasure_reward
+
+        self._action_mask = GW.make_action_mask(
+            [
+                GW.MOVE_UP,
+                GW.MOVE_RIGHT,
+                GW.MOVE_DOWN,
+                GW.MOVE_LEFT,
+            ],
+            self.num_agents,
+        )
+
+    def _generate_map(self, rng_key):
+        walls_key, decor_key, rng_key = jax.random.split(rng_key, 3)
+        noise = fractal_noise(
+            self.unpadded_width, self.unpadded_height, [2, 4, 5, 8, 10], walls_key
+        )
+
+        tiles = generate_decor_tiles(
+            self.unpadded_width, self.unpadded_height, decor_key
+        )
+        tiles = jnp.where(noise > 0.05, jnp.int8(GW.TILE_DESTRUCTIBLE_WALL), tiles)
+
+        # get the empty tiles for spawning
+        x_spawns, y_spawns = jnp.where(
+            tiles == GW.TILE_EMPTY,
+            size=self.unpadded_width * self.unpadded_height,
+            fill_value=jnp.int8(-1),
+        )
+        spawn_count = jnp.sum(tiles == GW.TILE_EMPTY)
+
+        # pad the tiles
+        tiles = jnp.pad(
+            tiles,
+            pad_width=(
+                (self.pad_width, self.pad_width),
+                (self.pad_height, self.pad_height),
+            ),
+            mode="constant",
+            constant_values=GW.TILE_WALL,
+        )
+
+        # pad the empty tiles
+        y_spawns = y_spawns + self.pad_height
+        x_spawns = x_spawns + self.pad_width
+        spawn_pos = jnp.stack((x_spawns, y_spawns), axis=1)
+
+        return tiles, spawn_pos, spawn_count
+
+    def reset(self, rng_key: jax.Array) -> tuple[ReturnDiggingState, TimeStep]:
+        map_key, pos_key = jax.random.split(rng_key)
+
+        map, spawn_pos, spawn_count = self._generate_map(map_key)
+
+        unpadded_map = map[
+            self.pad_width : -self.pad_width, self.pad_height : -self.pad_height
+        ]
+
+        pos_x, pos_y = choose_positions(
+            unpadded_map,
+            self.num_flags + self.num_agents,
+            pos_key,
+            replace=False,
+        )
+
+        pos_x = pos_x + self.pad_width
+        pos_y = pos_y + self.pad_height
+        positions = jnp.stack((pos_x, pos_y), axis=1)
+        flag_pos = positions[: self.num_flags]
+        agents_pos = positions[self.num_flags :]
+
+        if self._config.eval_map:
+            o = map
+            map = map.at[36:45, 10:30].set(GW.TILE_DESTRUCTIBLE_WALL)
+            map = map.at[42:45, 17:23].set(o[42:45, 17:23])
+            agents_pos = agents_pos.at[0].set([44, 22])
+            map = map.at[43, 18].set(GW.TILE_FLAG)
+        else:
+            map = map.at[flag_pos[:, 0], flag_pos[:, 1]].set(GW.TILE_FLAG)
+
+        state = ReturnDiggingState(
+            map=map,
+            spawn_pos=spawn_pos,
+            spawn_count=spawn_count,
+            agents_pos=agents_pos,
+            agents_timeout=jnp.zeros((self.num_agents,), dtype=jnp.int32),
+            found_reward=jnp.zeros((self.num_agents,), dtype=jnp.bool_),
+            time=jnp.int32(0),
+            rewards=jnp.float32(0.0),
+        )
+
+        actions = jnp.zeros((self.num_agents,), dtype=jnp.int32)
+        rewards = jnp.zeros((self.num_agents,), dtype=jnp.float32)
+
+        return state, self.encode_observations(state, actions, rewards)
+
+    def load_map(self, map: str):
+        tiles = np.zeros((self.unpadded_width, self.unpadded_height), dtype=np.int8)
+
+        x = 0
+        y = self.unpadded_height
+
+        agent_positions = []
+        spawn_positions = []
+
+        for c in map:
+            if c == "\n":
+                x = 0
+                y -= 1
+            else:
+                match x:
+                    case "x":
+                        tiles[x, y] = GW.TILE_DESTRUCTIBLE_WALL
+                    case "a":
+                        agent_positions.append(
+                            [self.pad_width + x, self.pad_height + y]
+                        )
+                    case "f":
+                        tiles[x, y] = GW.TILE_FLAG
+                    case _:
+                        spawn_positions.append(
+                            [self.pad_width + x, self.pad_height + y]
+                        )
+
+                x += 1
+
+        tiles = jnp.asarray(tiles)
+        # pad the tiles
+        tiles = jnp.pad(
+            tiles,
+            pad_width=(
+                (self.pad_width, self.pad_width),
+                (self.pad_height, self.pad_height),
+            ),
+            mode="constant",
+            constant_values=GW.TILE_WALL,
+        )
+
+        state = ReturnDiggingState(
+            map=tiles,
+            spawn_pos=jnp.array(spawn_positions, jnp.int32),
+            spawn_count=jnp.int32(len(spawn_positions)),
+            agents_pos=jnp.arange(agent_positions, jnp.int32),
+            agents_timeout=jnp.zeros((self.num_agents,), dtype=jnp.int32),
+            found_reward=jnp.zeros((self.num_agents,), dtype=jnp.bool_),
+            time=jnp.int32(0),
+            rewards=jnp.float32(0.0),
+        )
+
+        actions = jnp.zeros((self.num_agents,), dtype=jnp.int32)
+        rewards = jnp.zeros((self.num_agents,), dtype=jnp.float32)
+
+        return state, self.encode_observations(state, actions, rewards)
+
+    @cached_property
+    def observation_spec(self) -> ObservationSpec:
+        return GW.make_obs_spec(self.view_width, self.view_height)
+
+    @cached_property
+    def action_spec(self) -> DiscreteActionSpec:
+        return DiscreteActionSpec(num_actions=GW.NUM_ACTIONS)
+
+    @property
+    def is_jittable(self) -> bool:
+        return True
+
+    @property
+    def num_agents(self) -> int:
+        return self._num_agents
+
+    def step(
+        self, state: ReturnDiggingState, action: jax.Array, rng_key: jax.Array
+    ) -> tuple[ReturnDiggingState, TimeStep]:
+        @partial(jax.vmap, in_axes=(0, 0, 0, 0), out_axes=(0, 0, 0, 0))
+        def _step_agent(local_position, timeout, local_action, random_position):
+            def _step_timeout(local_position, timeout, local_action, random_position):
+                return local_position, local_position, timeout - 1, 0.0
+
+            def _step_move(local_position, timeout, local_action, random_position):
+                target_pos = local_position + GW.DIRECTIONS[local_action]
+
+                new_tile = state.map[target_pos[0], target_pos[1]]
+
+                # don't move if we are moving into a wall
+                new_pos = jnp.where(
+                    jnp.logical_or(
+                        new_tile == GW.TILE_WALL, new_tile == GW.TILE_DESTRUCTIBLE_WALL
+                    ),
+                    local_position,
+                    target_pos,
+                )
+
+                found_treasure = new_tile == GW.TILE_FLAG
+                reward = jnp.where(found_treasure, self.treasure_reward, 0.0)
+
+                # randomize position if the agent finds the reward
+                new_pos = jnp.where(found_treasure, random_position, new_pos)
+
+                # sets a timeout of the tile is dug
+                timeout = jnp.where(
+                    new_tile == GW.TILE_DESTRUCTIBLE_WALL, self.digging_timeout, 0
+                )
+
+                return new_pos, target_pos, timeout, reward
+
+            return jax.lax.cond(
+                timeout > 0,
+                _step_timeout,
+                _step_move,
+                local_position,
+                timeout,
+                local_action,
+                random_position,
+            )
+
+        random_positions = state.spawn_pos[
+            jax.random.randint(
+                rng_key, (self._num_agents,), minval=0, maxval=state.spawn_count
+            )
+        ]
+        new_position, target_pos, timeout, rewards = _step_agent(
+            state.agents_pos, state.agents_timeout, action, random_positions
+        )
+
+        # dig actions
+        target_tiles = state.map[target_pos[:, 0], target_pos[:, 1]]
+        map = state.map.at[target_pos[:, 0], target_pos[:, 1]].set(
+            jnp.where(
+                target_tiles == GW.TILE_DESTRUCTIBLE_WALL, GW.TILE_EMPTY, target_tiles
+            )
+        )
+        # /dig actions
+
+        state = state._replace(
+            agents_pos=new_position,
+            agents_timeout=timeout,
+            found_reward=jnp.logical_or(state.found_reward, rewards),
+            time=state.time + 1,
+            rewards=state.rewards + jnp.mean(rewards),
+            map=map,
+        )
+
+        return state, self.encode_observations(state, action, rewards)
+
+    def _render_tiles(self, state: ReturnDiggingState):
+        tiles = state.map
+        tiles = tiles.at[state.agents_pos[:, 0], state.agents_pos[:, 1]].set(
+            GW.AGENT_GENERIC
+        )
+
+        directions = jnp.zeros_like(tiles, dtype=jnp.int8)
+        teams = jnp.zeros_like(tiles, dtype=jnp.int8)
+        health = jnp.zeros_like(tiles, dtype=jnp.int8)
+
+        return jnp.concatenate(
+            (
+                tiles[..., None],
+                directions[..., None],
+                teams[..., None],
+                health[..., None],
+            ),
+            axis=-1,
+        )
+
+    def encode_observations(
+        self, state: ReturnDiggingState, actions, rewards
+    ) -> TimeStep:
+        @partial(jax.vmap, in_axes=(None, 0))
+        def _encode_view(tiles, positions):
+            return jax.lax.dynamic_slice(
+                tiles,
+                (
+                    positions[0] - self.view_width // 2,
+                    positions[1] - self.view_height // 2,
+                    0,
+                ),
+                (self.view_width, self.view_height, self.observation_spec.shape[-1]),
+            )
+
+        tiles = self._render_tiles(state)
+        view = _encode_view(tiles, state.agents_pos)
+
+        time = jnp.repeat(state.time[None], self.num_agents, axis=0)
+
+        return TimeStep(
+            obs=view,
+            time=time,
+            last_action=actions,
+            last_reward=rewards,
+            action_mask=self._action_mask,
+            terminated=jnp.equal(time, self._length - 1),
+        )
+
+    def create_placeholder_logs(self):
+        return {"rewards": jnp.float32(0.0)}
+
+    def create_logs(self, state: ReturnDiggingState):
+        return {"rewards": state.rewards}
+
+    def get_render_state(self, state: ReturnDiggingState) -> GridRenderState:
+        tiles = self._render_tiles(state)
+
+        return GridRenderState(
+            tilemap=tiles,
+            agent_positions=state.agents_pos,
+        )
+
+    def get_render_settings(self) -> GridRenderSettings:
+        return GridRenderSettings(
+            tile_width=self.unpadded_width,
+            tile_height=self.unpadded_height,
+            view_width=self.view_width,
+            view_height=self.view_height,
+        )