mapox 0.1.0__tar.gz

mapox-0.1.0/.gitignore

@@ -0,0 +1,10 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv

mapox-0.1.0/.python-version

@@ -0,0 +1 @@
+3.11

mapox-0.1.0/PKG-INFO

@@ -0,0 +1,17 @@
+Metadata-Version: 2.4
+Name: mapox
+Version: 0.1.0
+Summary: Multi-Agent Partially Observable gridworlds in JAX
+Author-email: Gabriel Keith <gfk1995@gmail.com>
+Keywords: gridworld,jax,multi-agent,pomdp,reinforcement-learning
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.11
+Requires-Dist: einops>=0.8.1
+Requires-Dist: jax>=0.7.2
+Requires-Dist: pydantic>=2.12.5
+Requires-Dist: pygame-ce>=2.5.6
+Requires-Dist: python-ffmpeg>=2.0.12

mapox-0.1.0/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-0.1.0/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-0.1.0/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-0.1.0/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-0.1.0/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)