JAXtari 0.1.0__py3-none-any.whl

jaxatari/__init__.py

@@ -0,0 +1,32 @@
+from pathlib import Path
+from platformdirs import user_data_dir
+
+# 1. Define the path (Must match the installer script exactly)
+# appname="jaxatari", appauthor="mycompany" (or whatever you used)
+DATA_DIR = Path(user_data_dir("jaxatari"))
+MARKER_FILE = DATA_DIR / ".ownership_confirmed"
+ALT_SPRITES_MARKER_FILE = DATA_DIR / ".alternative_sprites_installed"
+
+def check_ownership():
+    """
+    Verifies that the user has accepted the license and confirmed ownership
+    of the original hardware/software by looking for the marker file.
+    """
+    if not (MARKER_FILE.exists() or ALT_SPRITES_MARKER_FILE.exists()):
+        # Raise a clear, blocking error
+        raise RuntimeError(
+            "\n"
+            "❌  SPRITES NOT INSTALLED\n"
+            "----------------------------------------------------\n"
+            "JaxAtari needs sprite assets before environments can start.\n"
+            "You can either confirm your ownership of the original Atari 2600 ROMs and install sprites,\n"
+            "or continue with replacement/custom sprites.\n\n"
+            "Please run the following command in your terminal:\n\n"
+            "    .venv/bin/install-sprites\n"
+            "    or\n"
+            "    python3 scripts/install_sprites.py\n"
+            "----------------------------------------------------\n"
+        )
+
+# ... rest of your package imports ...
+from jaxatari.core import make, list_available_games

jaxatari/core.py

@@ -0,0 +1,207 @@
+import importlib
+import inspect
+import warnings
+
+from jaxatari.environment import JaxEnvironment
+from jaxatari.renderers import JAXGameRenderer
+from jaxatari.modification import apply_modifications
+from jaxatari.wrappers import JaxatariWrapper
+
+from . import check_ownership
+
+
+def _warn_deprecated_obs_to_flat_array(env: JaxEnvironment) -> None:
+    """Warn if legacy obs_to_flat_array is present on the environment."""
+    if hasattr(env, "obs_to_flat_array") and callable(getattr(env, "obs_to_flat_array")):
+        warnings.warn(
+            "Environment exposes deprecated obs_to_flat_array(). "
+            "Observations should now be flax.struct.dataclasses using ObjectObservation "
+            "for objects or plain arrays for observations like lives, score, etc. "
+            "Depending on legacy obs_to_flat_array might lead to unforseen issues with wrappers.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
+
+
+# Map of game names to their module paths (commented out games are WIP and will be supported in the near future)
+GAME_MODULES = {
+    "amidar": "jaxatari.games.jax_amidar",
+    "airraid": "jaxatari.games.jax_airraid",
+    "alien": "jaxatari.games.jax_alien",
+    "asterix": "jaxatari.games.jax_asterix",
+    "asteroids": "jaxatari.games.jax_asteroids",
+    "atlantis": "jaxatari.games.jax_atlantis",
+    "bankheist": "jaxatari.games.jax_bankheist",
+    "beamrider": "jaxatari.games.jax_beamrider",
+    "berzerk": "jaxatari.games.jax_berzerk",
+    "blackjack": "jaxatari.games.jax_blackjack",
+    "breakout": "jaxatari.games.jax_breakout",
+    "casino": "jaxatari.games.jax_casino",
+    "centipede": "jaxatari.games.jax_centipede",
+    "choppercommand": "jaxatari.games.jax_choppercommand",
+    "donkeykong": "jaxatari.games.jax_donkeykong",
+    "enduro": "jaxatari.games.jax_enduro",
+    "fishingderby": "jaxatari.games.jax_fishingderby",
+    "flagcapture": "jaxatari.games.jax_flagcapture",
+    "freeway": "jaxatari.games.jax_freeway",
+    "frostbite": "jaxatari.games.jax_frostbite",
+    "galaxian": "jaxatari.games.jax_galaxian",
+    "gravitar": "jaxatari.games.jax_gravitar",
+    "hangman": "jaxatari.games.jax_hangman",
+    "hauntedhouse": "jaxatari.games.jax_hauntedhouse",
+    "humancannonball": "jaxatari.games.jax_humancannonball",
+    "kangaroo": "jaxatari.games.jax_kangaroo",
+    "kingkong": "jaxatari.games.jax_kingkong",
+    "klax": "jaxatari.games.jax_klax",
+    "lasergates": "jaxatari.games.jax_lasergates",
+    "namethisgame": "jaxatari.games.jax_namethisgame",
+    "phoenix": "jaxatari.games.jax_phoenix",
+    "pong": "jaxatari.games.jax_pong",
+    "qbert": "jaxatari.games.jax_qbert",
+    "riverraid": "jaxatari.games.jax_riverraid",
+    "seaquest": "jaxatari.games.jax_seaquest",
+    "sirlancelot": "jaxatari.games.jax_sirlancelot",
+    "skiing": "jaxatari.games.jax_skiing",
+    "slotmachine": "jaxatari.games.jax_slotmachine",
+    "spaceinvaders": "jaxatari.games.jax_spaceinvaders",
+    "spacewar": "jaxatari.games.jax_spacewar",
+    "surround": "jaxatari.games.jax_surround",
+    "tennis": "jaxatari.games.jax_tennis",
+    "tetris": "jaxatari.games.jax_tetris",
+    "timepilot": "jaxatari.games.jax_timepilot",
+    "tron": "jaxatari.games.jax_tron",
+    "turmoil": "jaxatari.games.jax_turmoil",
+    "venture": "jaxatari.games.jax_venture",
+    "videocheckers": "jaxatari.games.jax_videocheckers",
+    "videocube": "jaxatari.games.jax_videocube",
+    "videopinball": "jaxatari.games.jax_videopinball",
+    "wordzapper": "jaxatari.games.jax_wordzapper",
+    "mspacman": "jaxatari.games.jax_mspacman",
+    "montezumarevenge": "jaxatari.games.jax_montezumarevenge",
+    "pacman": "jaxatari.games.jax_pacman",
+}
+
+# Mod modules registry: for each game, provide the Controller class path
+MOD_MODULES = {
+    "pong": "jaxatari.games.mods.pong_mods.PongEnvMod",
+    "kangaroo": "jaxatari.games.mods.kangaroo_mods.KangarooEnvMod",
+    "freeway": "jaxatari.games.mods.freeway_mods.FreewayEnvMod",
+    "breakout": "jaxatari.games.mods.breakout_mods.BreakoutEnvMod",
+    "seaquest": "jaxatari.games.mods.seaquest_mods.SeaquestEnvMod",
+    "videopinball": "jaxatari.games.mods.videopinball_mods.VideoPinballEnvMod",
+    'tennis': "jaxatari.games.mods.tennis_mods.TennisEnvMod",
+    "fishingderby": "jaxatari.games.mods.fishingderby_mods.FishingDerbyEnvMod",
+    "atlantis": "jaxatari.games.mods.atlantis_mods.AtlantisEnvMod",
+    "bankheist": "jaxatari.games.mods.bankheist_mods.BankHeistEnvMod",
+    "montezumarevenge": "jaxatari.games.mods.montezuma_revenge_mods.MontezumaRevengeEnvMod",
+    "frostbite": "jaxatari.games.mods.frostbite_mods.FrostbiteEnvMod",
+    "gravitar": "jaxatari.games.mods.gravitar_mods.GravitarEnvMod",
+    "phoenix": "jaxatari.games.mods.phoenix_mods.PhoenixEnvMod",
+    "enduro": "jaxatari.games.mods.enduro_mods.EnduroEnvMod",
+    "qbert": "jaxatari.games.mods.qbert_mods.QbertEnvMod",
+    "mspacman": "jaxatari.games.mods.mspacman_mods.MsPacmanEnvMod",
+    "beamrider": "jaxatari.games.mods.beamrider_mods.BeamRiderEnvMod",
+    "venture": "jaxatari.games.mods.venture_mods.VentureEnvMod",
+    "spaceinvaders": "jaxatari.games.mods.spaceinvaders_mods.SpaceInvadersEnvMod",
+    "skiing": "jaxatari.games.mods.skiing_mods.SkiingEnvMod",
+    "alien": "jaxatari.games.mods.alien_mods.AlienEnvMod",
+    "asteroids": "jaxatari.games.mods.asteroids_mods.AsteroidsEnvMod"
+}
+
+
+def list_available_games() -> list[str]:
+    """Lists all available, registered games."""
+    return list(GAME_MODULES.keys())
+
+
+def make(game_name: str, 
+         mode: int = 0, 
+         difficulty: int = 0,
+         mods_config: list = None, # deprecated, output warning if its used
+         mods: list = None,
+         allow_conflicts: bool = False
+         ) -> JaxEnvironment:
+    """
+    Creates and returns a JaxAtari game environment instance.
+    This is the main entry point for creating environments.
+
+    If 'mods' is provided, this function applies the
+    full two-stage modding pipeline:
+    1. Pre-scans for constant overrides.
+    2. Instantiates the base env with modded constants.
+    3. Applies the internal 'JaxAtariModController'.
+    4. Wraps the env with the 'JaxAtariModWrapper'.
+
+    Args:
+        game_name: Name of the game to load (e.g., "pong").
+        mode: Game mode.
+        difficulty: Game difficulty.
+        mods: List of modifications to apply (default: None).
+        allow_conflicts: Whether to allow conflicting mods (default: False).
+    Returns:
+        An instance of the specified game environment.
+    """
+
+    check_ownership()  # Ensure ownership confirmed
+
+    if isinstance(game_name, str):
+        game_name_clean = game_name.lower().replace("_", "").replace("-", "")
+        for key in GAME_MODULES:
+            if key.lower().replace("_", "").replace("-", "") == game_name_clean:
+                game_name = key
+                break
+
+    if mods_config is not None:
+        warnings.warn(
+            "'mods_config' is deprecated and will be removed in future versions. "
+            "Please use 'mods' instead.",
+            DeprecationWarning
+        )
+        mods = mods_config
+
+    if game_name not in GAME_MODULES:
+        raise NotImplementedError(
+            f"The game '{game_name}' does not exist. Available games: {list_available_games()}"
+        )
+    
+    try:
+        # 1. Load the base environment class
+        module = importlib.import_module(GAME_MODULES[game_name])
+        env_class = None
+        for _, obj in inspect.getmembers(module):
+            if inspect.isclass(obj) and issubclass(obj, JaxEnvironment) and obj is not JaxEnvironment:
+                env_class = obj
+                break
+        if env_class is None:
+            raise ImportError(f"No JaxEnvironment subclass found in {GAME_MODULES[game_name]}")
+
+        # 2. Mods need default consts for pre-scan; otherwise a single env_class() is enough.
+        if mods:
+            try:
+                base_consts = env_class().consts
+                env = apply_modifications(
+                    game_name=game_name,
+                    mods_config=mods,
+                    allow_conflicts=allow_conflicts,
+                    base_consts=base_consts,
+                    env_class=env_class,
+                    MOD_MODULES=MOD_MODULES
+                )
+                _warn_deprecated_obs_to_flat_array(env)
+                return env
+            except NotImplementedError as e:
+                # Mod module not defined for this game - fall back to base environment
+                warnings.warn(
+                    f"Mods requested for '{game_name}' but no mod module is available. "
+                    f"Creating base environment without mods. Error: {e}",
+                    UserWarning
+                )
+
+        env = env_class()
+        _warn_deprecated_obs_to_flat_array(env)
+        return env
+
+    except (ImportError, NotImplementedError) as e:
+        # Only wrap registration/import errors - let intentional errors (ValueError, etc.) propagate
+        raise ImportError(f"Failed to load game '{game_name}': {e}") from e

jaxatari/environment.py

@@ -0,0 +1,289 @@
+from enum import Enum
+from typing import Tuple, Generic, TypeVar
+import jax.numpy as jnp
+import jax.random as jrandom
+import warnings
+from jaxatari.spaces import Space
+from flax import struct
+
+EnvObs = TypeVar("EnvObs")
+EnvState = TypeVar("EnvState")
+EnvInfo = TypeVar("EnvInfo")
+EnvConstants = TypeVar("EnvConstants")
+
+class JAXAtariAction:
+    """
+    "Namespace" for Atari action integer constants.
+    These are directly usable in JAX arrays.
+    """
+    NOOP: int = 0
+    FIRE: int = 1
+    UP: int = 2
+    RIGHT: int = 3
+    LEFT: int = 4
+    DOWN: int = 5
+    UPRIGHT: int = 6
+    UPLEFT: int = 7
+    DOWNRIGHT: int = 8
+    DOWNLEFT: int = 9
+    UPFIRE: int = 10
+    RIGHTFIRE: int = 11
+    LEFTFIRE: int = 12
+    DOWNFIRE: int = 13
+    UPRIGHTFIRE: int = 14
+    UPLEFTFIRE: int = 15
+    DOWNRIGHTFIRE: int = 16
+    DOWNLEFTFIRE: int = 17
+
+    @classmethod
+    def get_all_values(cls) -> jnp.ndarray:
+        # For fixed action sets, explicit listing is safest and clearest.
+        return jnp.array([
+            cls.NOOP, cls.FIRE, cls.UP, cls.RIGHT, cls.LEFT, cls.DOWN,
+            cls.UPRIGHT, cls.UPLEFT, cls.DOWNRIGHT, cls.DOWNLEFT,
+            cls.UPFIRE, cls.RIGHTFIRE, cls.LEFTFIRE, cls.DOWNFIRE,
+            cls.UPRIGHTFIRE, cls.UPLEFTFIRE, cls.DOWNRIGHTFIRE, cls.DOWNLEFTFIRE
+        ], dtype=jnp.int32)
+
+@struct.dataclass
+class ObjectObservation:
+    """
+    Dataclass for object centric observations of objects in jaxatari environments. 
+    Can hold 1 to N objects of the same type (for example 12 sharks in seaquest or 1 player ship in asteroids).
+    Should always be instantiated via the create() classmethod to ensure proper default handling.
+    Attributes:
+        x: x position of the object.
+        y: y position of the object.
+        width: width of the object.
+        height: height of the object.
+        active: whether the object is currently active.
+    """
+    x: jnp.ndarray          # obligatory (int8)
+    y: jnp.ndarray          # obligatory (int8)
+    width: jnp.ndarray      # obligatory (int8)
+    height: jnp.ndarray     # obligatory (int8)
+
+    # --- Additional attributes (will be set to 0 if not used) ---
+    active: jnp.ndarray = struct.field(default_factory=lambda: jnp.array(1))          # whether the object is currently active (0 or 1)
+    visual_id: jnp.ndarray = struct.field(default_factory=lambda: jnp.array(0))   # visual identifier of the object (different color sprites for example)
+    state: jnp.ndarray = struct.field(default_factory=lambda: jnp.array(0))       # state of the object, for example is the ghost in pacman vulnerable [blinking] or not [static] (format depends on game, see the game docs)
+    orientation: jnp.ndarray = struct.field(default_factory=lambda: jnp.array(0))       # angle of the object (format depends on game, see the game docs)
+
+    @classmethod
+    def create(cls, x, y, width, height, active=None, visual_id=None, state=None, orientation=None):
+        # Helper to handle defaults
+        if active is None: active = jnp.ones_like(x, dtype=jnp.int32)
+        if visual_id is None: visual_id = jnp.zeros_like(x, dtype=jnp.int32)
+        if state is None: state = jnp.zeros_like(x, dtype=jnp.int32)
+        if orientation is None: orientation = jnp.zeros_like(x, dtype=jnp.int32)
+        return cls(x=x, y=y, width=width, height=height, active=active, visual_id=visual_id, state=state, orientation=orientation)
+
+    def __repr__(self):
+        try:
+            # Handle scalar case (0-d arrays)
+            if self.x.ndim == 0:
+                try:
+                    # Try to get concrete values for cleaner output
+                    x, y = int(self.x), int(self.y)
+                    w, h = int(self.width), int(self.height)
+                    act = int(self.active)
+                    ori = float(self.orientation)
+                    st = int(self.state)
+                    vid = int(self.visual_id)
+                    status = "ACTIVE" if act else "INACTIVE"
+                    return (f"Object(Single, {status}): Pos=({x}, {y}) | Size=({w}, {h}) | "
+                            f"Ori={ori:.1f} | State={st} | VisID={vid}")
+                except:
+                    # Fallback for Tracers
+                    return f"Object(Single): Pos=({self.x}, {self.y}) | Active={self.active}"
+            
+            # Handle vector case (1-d arrays)
+            n = self.x.shape[0]
+            lines = [f"ObjectGroup(count={n}):"]
+            
+            # Limit print length if too huge
+            limit = min(n, 20) 
+            
+            for i in range(limit):
+                try:
+                    # Try to extract concrete values
+                    act = int(self.active[i])
+                    status = "ACTIVE" if act else "   -  " # Dim inactive ones
+                    
+                    x, y = int(self.x[i]), int(self.y[i])
+                    w, h = int(self.width[i]), int(self.height[i])
+                    ori = float(self.orientation[i])
+                    st = int(self.state[i])
+                    vid = int(self.visual_id[i])
+                    
+                    # Formatted table row
+                    line = (f"  [{i:2d}] {status} | Pos: ({x:3d}, {y:3d}) | Size: ({w:2d}, {h:2d}) | "
+                            f"Ori: {ori:5.1f} | State: {st:2d} | VisID: {vid:2d}")
+                except:
+                    # Fallback for Tracers
+                    line = f"  [{i}] Active={self.active[i]} | Pos=({self.x[i]}, {self.y[i]})"
+                
+                lines.append(line)
+            
+            if n > limit:
+                lines.append(f"  ... ({n - limit} more objects) ...")
+                
+            return "\n".join(lines)
+        except Exception as e:
+            return f"ObjectObservation(Error in __repr__: {e})"
+
+
+class JaxEnvironment(Generic[EnvState, EnvObs, EnvInfo, EnvConstants]):
+    """
+    Abstract class for a JAX environment.
+    Generics:
+    EnvState: The type of the environment state.
+    EnvObs: The type of the observation.
+    EnvInfo: The type of the additional information.
+    EnvConstants: The type of the environment constants.
+    """
+
+    def __init__(self, consts: EnvConstants = None):
+        if consts is not None:
+            # Check for legacy NamedTuple usage (has _fields but is not a PyTreeNode)
+            is_named_tuple = isinstance(consts, tuple) and hasattr(consts, '_fields')
+            # Check if it's a Flax PyTreeNode (flax.struct.dataclass instances)
+            try:
+                from flax import struct
+                is_flax_node = isinstance(consts, struct.PyTreeNode)
+            except (ImportError, AttributeError):
+                is_flax_node = False
+
+            if is_named_tuple and not is_flax_node:
+                warnings.warn(
+                    f"Performance Warning: {self.__class__.__name__}.consts is a 'NamedTuple'. "
+                    "This prevents JAX from treating constants as static metadata, potentially causing excessive recompilation. "
+                    "Future versions will require 'flax.struct.PyTreeNode' (and the states/observations/info to flax.struct.dataclass/PyTreeNode). "
+                    "Please refactor your constants class.",
+                    UserWarning,
+                    stacklevel=2
+                )
+
+        self.consts = consts
+
+        # --- MODDING INFRASTRUCTURE ---
+        # Functional: Tracks which renderer methods mods have patched.
+        # Used by wrappers to safely transfer patches during renderer swaps.
+        self._patched_renderer_methods = []
+
+        # Functional: Explicit registry of jitted callables that must be invalidated
+        # when renderer hot-swaps occur (e.g., native downscaling).
+        self._jit_invalidation_targets = []
+        # Functional: mutation epoch + tripwire controls for detecting risky
+        # post-trace monkeypatching.
+        self._jit_mutation_epoch = 0
+        self._jit_tripwire_enabled = True
+        
+        # Informational: Structured audit log of every change made by the mod system.
+        # Machine-parseable: dict of category -> set of names that were changed.
+        # Categories: "attribute", "method", "constant", "asset".
+        self._mod_history = {
+            "attribute": set(),
+            "method": set(),
+            "constant": set(),
+            "asset": set(),
+        }
+
+    def reset(self, key: jrandom.PRNGKey=None) -> Tuple[EnvObs, EnvState]:
+        """
+        Resets the environment to the initial state.
+        Returns: The initial observation and the initial environment state.
+
+        """
+        raise NotImplementedError("Abstract method")
+
+    def step(
+        self, state: EnvState, action
+    ) -> Tuple[EnvObs, EnvState, float, bool, EnvInfo]:
+        """
+        Takes a step in the environment.
+        Args:
+            state: The current environment state.
+            action: The action to take.
+
+        Returns: The observation, the new environment state, the reward, whether the state is terminal, and additional info.
+
+        """
+        raise NotImplementedError("Abstract method")
+
+    def render(self, state: EnvState) -> Tuple[jnp.ndarray]:
+        """
+        Renders the environment state to a single image.
+        Args:
+            state: The environment state.
+
+        Returns: A single image of the environment state.
+
+        """
+        raise NotImplementedError("Abstract method")
+
+    def action_space(self) -> Space:
+        """
+        Returns the action space of the environment as an array containing the actions that can be taken.
+        Returns: The action space of the environment as an array.
+        """
+        raise NotImplementedError("Abstract method")
+
+    def observation_space(self) -> Space:
+        """
+        Returns the observation space of the environment.
+        Returns: The observation space of the environment.
+        """
+        raise NotImplementedError("Abstract method")
+    
+    def image_space(self) -> Space:
+        """
+        Returns the image space of the environment.
+        Returns: The image space of the environment.
+        """
+        raise NotImplementedError("Abstract method")
+
+    def _get_observation(self, state: EnvState) -> EnvObs:
+        """
+        Converts the environment state to the observation by filtering out non-relevant information.
+        Args:
+            state: The environment state.
+
+        Returns: observation
+
+        """
+        raise NotImplementedError("Abstract method")
+
+    def _get_info(self, state: EnvState, all_rewards: jnp.array = None) -> EnvInfo:
+        """
+        Extracts information from the environment state that is not relevant for the agent.
+        Args:
+            state: The environment state.
+
+        Returns: info
+
+        """
+        raise NotImplementedError("Abstract method")
+
+    def _get_reward(self, previous_state: EnvState, state: EnvState) -> float:
+        """
+        Calculates the reward from the environment state.
+        Args:
+            previous_state: The previous environment state.
+            state: The environment state.
+
+        Returns: reward
+
+        """
+        raise NotImplementedError("Abstract method")
+
+    def _get_done(self, state: EnvState) -> bool:
+        """
+        Determines if the environment state is a terminal state
+        Args:
+            state: The environment state.
+
+        Returns: True if the state is terminal, False otherwise.
+
+        """
+        raise NotImplementedError("Abstract method")

jaxatari/games/amidar_mazes.py

@@ -0,0 +1,23 @@
+# Original amidar maze 
+import jax.numpy as jnp
+
+WIDTH = 160
+HEIGHT = 210
+PATH_THICKNESS_HORIZONTAL = 5
+PATH_THICKNESS_VERTICAL = 4
+MAX_ENEMIES = 6
+
+PATH_CORNERS = jnp.array([[16, 14], [40, 14], [56, 14], [72, 14], [84, 14], [100, 14], [116, 14], [140, 14], [16, 44], [32, 44], [40, 44], [52, 44], [56, 44], [64, 44], [72, 44], [84, 44], [92, 44], [100, 44], [104, 44], [116, 44], [124, 44], [140, 44], [16, 74], [28, 74], [32, 74], [52, 74], [60, 74], [64, 74], [92, 74], [96, 74], [104, 74], [124, 74], [128, 74], [140, 74], [16, 104], [28, 104], [36, 104], [60, 104], [72, 104], [84, 104], [96, 104], [120, 104], [128, 104], [140, 104], [16, 134], [36, 134], [40, 134], [64, 134], [72, 134], [84, 134], [92, 134], [116, 134], [120, 134], [140, 134], [16, 164], [40, 164], [64, 164], [92, 164], [116, 164], [140, 164]], dtype=jnp.int32)
+HORIZONTAL_PATH_EDGES = jnp.array([[[16, 14], [40, 14]], [[16, 44], [32, 44]], [[16, 74], [28, 74]], [[16, 104], [28, 104]], [[16, 134], [36, 134]], [[16, 164], [40, 164]], [[28, 74], [32, 74]], [[28, 104], [36, 104]], [[32, 44], [40, 44]], [[32, 74], [52, 74]], [[36, 104], [60, 104]], [[36, 134], [40, 134]], [[40, 14], [56, 14]], [[40, 44], [52, 44]], [[40, 134], [64, 134]], [[40, 164], [64, 164]], [[52, 44], [56, 44]], [[52, 74], [60, 74]], [[56, 14], [72, 14]], [[56, 44], [64, 44]], [[60, 74], [64, 74]], [[60, 104], [72, 104]], [[64, 44], [72, 44]], [[64, 74], [92, 74]], [[64, 134], [72, 134]], [[64, 164], [92, 164]], [[72, 14], [84, 14]], [[72, 44], [84, 44]], [[72, 104], [84, 104]], [[72, 134], [84, 134]], [[84, 14], [100, 14]], [[84, 44], [92, 44]], [[84, 104], [96, 104]], [[84, 134], [92, 134]], [[92, 44], [100, 44]], [[92, 74], [96, 74]], [[92, 134], [116, 134]], [[92, 164], [116, 164]], [[96, 74], [104, 74]], [[96, 104], [120, 104]], [[100, 14], [116, 14]], [[100, 44], [104, 44]], [[104, 44], [116, 44]], [[104, 74], [124, 74]], [[116, 14], [140, 14]], [[116, 44], [124, 44]], [[116, 134], [120, 134]], [[116, 164], [140, 164]], [[120, 104], [128, 104]], [[120, 134], [140, 134]], [[124, 44], [140, 44]], [[124, 74], [128, 74]], [[128, 74], [140, 74]], [[128, 104], [140, 104]]], dtype=jnp.int32)
+VERTICAL_PATH_EDGES = jnp.array([[[16, 14], [16, 44]], [[16, 44], [16, 74]], [[16, 74], [16, 104]], [[16, 104], [16, 134]], [[16, 134], [16, 164]], [[28, 74], [28, 104]], [[32, 44], [32, 74]], [[36, 104], [36, 134]], [[40, 14], [40, 44]], [[40, 134], [40, 164]], [[52, 44], [52, 74]], [[56, 14], [56, 44]], [[60, 74], [60, 104]], [[64, 44], [64, 74]], [[64, 134], [64, 164]], [[72, 14], [72, 44]], [[72, 104], [72, 134]], [[84, 14], [84, 44]], [[84, 104], [84, 134]], [[92, 44], [92, 74]], [[92, 134], [92, 164]], [[96, 74], [96, 104]], [[100, 14], [100, 44]], [[104, 44], [104, 74]], [[116, 14], [116, 44]], [[116, 134], [116, 164]], [[120, 104], [120, 134]], [[124, 44], [124, 74]], [[128, 74], [128, 104]], [[140, 14], [140, 44]], [[140, 44], [140, 74]], [[140, 74], [140, 104]], [[140, 104], [140, 134]], [[140, 134], [140, 164]]], dtype=jnp.int32)
+PATH_EDGES = jnp.concatenate((HORIZONTAL_PATH_EDGES, VERTICAL_PATH_EDGES), axis=0)
+RECTANGLES = jnp.array([[1, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0], [0, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0], [0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 1, 1, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0], [0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0], [0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1]], dtype=jnp.int32)
+RECTANGLE_BOUNDS = jnp.array([[16, 14, 40, 44], [40, 14, 56, 44], [56, 14, 72, 44], [72, 14, 84, 44], [84, 14, 100, 44], [100, 14, 116, 44], [116, 14, 140, 44], [16, 44, 32, 74], [32, 44, 52, 74], [52, 44, 64, 74], [64, 44, 92, 74], [92, 44, 104, 74], [104, 44, 124, 74], [124, 44, 140, 74], [16, 74, 28, 104], [28, 74, 60, 104], [60, 74, 96, 104], [96, 74, 128, 104], [128, 74, 140, 104], [16, 104, 36, 134], [36, 104, 72, 134], [72, 104, 84, 134], [84, 104, 120, 134], [120, 104, 140, 134], [16, 134, 40, 164], [40, 134, 64, 164], [64, 134, 92, 164], [92, 134, 116, 164], [116, 134, 140, 164]], dtype=jnp.int32)
+CORNER_RECTANGLES = jnp.array([0, 6, 24, 28], dtype=jnp.int32)
+SHORT_PATHS = jnp.array([[23, 24, 6], [45, 46, 11], [11, 12, 16], [26, 27, 20], [28, 29, 35], [17, 18, 41], [51, 52, 46], [31, 32, 51]], dtype=jnp.int32)
+
+INITIAL_PLAYER_POSITION = jnp.array([140, 89], dtype=jnp.int32)
+INITIAL_ENEMY_POSITIONS = jnp.array([[16, 14], [16, 14], [44, 14], [16, 137], [52, 164], [16, 164]], dtype=jnp.int32)
+PLAYER_STARTING_PATH = jnp.array([85], dtype=jnp.int32)
+
+# Import these from this file where needed instead of modifying jax_amidar.py.