kaggle-environments 0.2.0__py3-none-any.whl → 1.20.0__py3-none-any.whl

kaggle_environments/envs/lux_ai_s3/luxai_s3/globals.py

@@ -0,0 +1,9 @@
+import os
+
+TERM_COLORS = True
+try:
+    TERM_COLORS = os.environ["LUX_COLORS"] == "False" if "LUX_COLORS" in os.environ else True
+except:
+    TERM_COLORS = False
+    # print("termcolor not installed, skipping dependency")
+    pass

kaggle_environments/envs/lux_ai_s3/luxai_s3/params.py

@@ -0,0 +1,101 @@
+from flax import struct
+
+MAP_TYPES = ["dev0", "random"]
+
+
+@struct.dataclass
+class EnvParams:
+    max_steps_in_match: int = 100
+    map_type: int = 1
+    """Map generation algorithm. Can change between games"""
+    map_width: int = 24
+    map_height: int = 24
+    num_teams: int = 2
+    match_count_per_episode: int = 5
+    """number of matches to play in one episode"""
+
+    # configs for units
+    max_units: int = 16
+    init_unit_energy: int = 100
+    min_unit_energy: int = 0
+    max_unit_energy: int = 400
+    unit_move_cost: int = 2
+    spawn_rate: int = 3
+
+    unit_sap_cost: int = 10
+    """
+    The unit sap cost is the amount of energy a unit uses when it saps another unit. Can change between games.
+    """
+    unit_sap_range: int = 4
+    """
+    The unit sap range is the range of the unit's sap action.
+    """
+    unit_sap_dropoff_factor: float = 0.5
+    """
+    The unit sap dropoff factor multiplied by unit_sap_drain
+    """
+    unit_energy_void_factor: float = 0.125
+    """
+    The unit energy void factor multiplied by unit_energy
+    """
+
+    # configs for energy nodes
+    max_energy_nodes: int = 6
+    max_energy_per_tile: int = 20
+    min_energy_per_tile: int = -20
+
+    max_relic_nodes: int = 6
+    """max relic nodes in the entire map. This number should be tuned carefully as relic node spawning code is hardcoded against this number 6"""
+    relic_config_size: int = 5
+    fog_of_war: bool = True
+    """
+    whether there is fog of war or not
+    """
+    unit_sensor_range: int = 2
+    """
+    The unit sensor range is the range of the unit's sensor.
+    Units provide "vision power" over tiles in range, equal to manhattan distance to the unit.
+
+    vision power > 0 that team can see the tiles properties
+    """
+
+    # nebula tile params
+    nebula_tile_vision_reduction: int = 1
+    """
+    The nebula tile vision reduction is the amount of vision reduction a nebula tile provides.
+    A tile can be seen if the vision power over it is > 0.
+    """
+
+    nebula_tile_energy_reduction: int = 0
+    """amount of energy nebula tiles reduce from a unit"""
+
+    nebula_tile_drift_speed: float = -0.05
+    """
+    how fast nebula tiles drift in one of the diagonal directions over time. If positive, flows to the top/right, negative flows to bottom/left
+    """
+    # TODO (stao): allow other kinds of symmetric drifts?
+
+    energy_node_drift_speed: int = 0.02
+    """
+    how fast energy nodes will move around over time
+    """
+    energy_node_drift_magnitude: int = 5
+
+    # option to change sap configurations
+
+
+env_params_ranges = dict(
+    # map_type=[1],
+    unit_move_cost=list(range(1, 6)),
+    unit_sensor_range=[1, 2, 3, 4],
+    nebula_tile_vision_reduction=list(range(0, 8)),
+    nebula_tile_energy_reduction=[0, 1, 2, 3, 5, 25],
+    unit_sap_cost=list(range(30, 51)),
+    unit_sap_range=list(range(3, 8)),
+    unit_sap_dropoff_factor=[0.25, 0.5, 1],
+    unit_energy_void_factor=[0.0625, 0.125, 0.25, 0.375],
+    # map randomizations
+    nebula_tile_drift_speed=[-0.15, -0.1, -0.05, -0.025, 0.025, 0.05, 0.1, 0.15],
+    energy_node_drift_speed=[0.01, 0.02, 0.03, 0.04, 0.05],
+    energy_node_drift_magnitude=list(range(3, 6)),
+)

kaggle_environments/envs/lux_ai_s3/luxai_s3/profiler.py

@@ -0,0 +1,141 @@
+import os
+import time
+from collections import defaultdict
+from typing import Literal
+
+import numpy as np
+import psutil
+import pynvml
+
+
+def flatten_dict_keys(d: dict, prefix=""):
+    """Flatten a dict by expanding its keys recursively."""
+    out = dict()
+    for k, v in d.items():
+        if isinstance(v, dict):
+            out.update(flatten_dict_keys(v, prefix + k + "/"))
+        else:
+            out[prefix + k] = v
+    return out
+
+
+class Profiler:
+    """
+    A simple class to help profile/benchmark simulator code
+    """
+
+    def __init__(self, output_format: Literal["stdout", "json"], synchronize_torch: bool = True) -> None:
+        self.output_format = output_format
+        self.synchronize_torch = synchronize_torch
+        self.stats = defaultdict(list)
+        # Initialize NVML
+        pynvml.nvmlInit()
+
+        # Get handle for the first GPU (index 0)
+        self.handle = pynvml.nvmlDeviceGetHandleByIndex(0)
+
+        # Get the PID of the current process
+        self.current_pid = os.getpid()
+
+    def log(self, msg):
+        """log a message to stdout"""
+        if self.output_format == "stdout":
+            print(msg)
+
+    def update_csv(self, csv_path: str, data: dict):
+        """Update a csv file with the given data (a dict representing a unique identifier of the result row)
+        and stats. If the file does not exist, it will be created. The update will replace an existing row
+        if the given data matches the data in the row. If there are multiple matches, only the first match
+        will be replaced and the rest are deleted"""
+        import os
+
+        import pandas as pd
+
+        if os.path.exists(csv_path):
+            df = pd.read_csv(csv_path)
+        else:
+            df = pd.DataFrame()
+        stats_flat = flatten_dict_keys(self.stats)
+        cond = None
+
+        for k in stats_flat:
+            if k not in df:
+                df[k] = None
+        for k in data:
+            if k not in df:
+                df[k] = None
+
+            mask = df[k].isna() if data[k] is None else df[k] == data[k]
+            if cond is None:
+                cond = mask
+            else:
+                cond = cond & mask
+        data_dict = {**data, **stats_flat}
+        if not cond.any():
+            df = pd.concat([df, pd.DataFrame(data_dict, index=[len(df)])])
+        else:
+            # replace the first instance
+            df.loc[df.loc[cond].index[0]] = data_dict
+            df.drop(df.loc[cond].index[1:], inplace=True)
+            # delete other instances
+        df.to_csv(csv_path, index=False)
+
+    def profile(self, function, name: str, total_steps: int, num_envs: int, trials=1):
+        print(f"start recording {name} metrics")
+        process = psutil.Process(os.getpid())
+        cpu_mem_use = process.memory_info().rss
+        gpu_mem_use = self.get_current_process_gpu_memory()
+        if gpu_mem_use is None:
+            gpu_mem_use = 0
+
+        for trial in range(trials):
+            stime = time.time()
+            function()
+            dt = time.time() - stime
+            # dt: delta time (s)
+            # fps: frames per second
+            # psps: parallel steps per second (number of env.step calls per second)
+            self.stats[name].append(
+                dict(
+                    dt=dt,
+                    fps=total_steps * num_envs / dt,
+                    psps=total_steps / dt,
+                    total_steps=total_steps,
+                    cpu_mem_use=cpu_mem_use,
+                    gpu_mem_use=gpu_mem_use,
+                )
+            )
+        # torch.cuda.synchronize()
+
+    def log_stats(self, name: str):
+        stats = self.stats[name]
+        more_than_one_trial = len(stats) > 1
+        if len(stats) == 0:
+            return
+        # average the stats
+        avg_stats = defaultdict(list)
+        for data in stats:
+            for k, v in data.items():
+                avg_stats[k].append(v)
+        stats = {k: {"avg": np.mean(v), "std": np.std(v) if len(v) > 1 else None} for k, v in avg_stats.items()}
+        self.log(f"{name} ({len(self.stats[name])} trials)")
+        self.log(
+            f"AVG: {stats['fps']['avg']:0.3f} steps/s, {stats['psps']['avg']:0.3f} parallel steps/s, {stats['total_steps']['avg']} steps in {stats['dt']['avg']:0.3f}s"
+        )
+        if more_than_one_trial:
+            self.log(
+                f"STD: {stats['fps']['std']:0.3f} steps/s, {stats['psps']['std']:0.3f} parallel steps/s, {stats['total_steps']['std']} steps in {stats['dt']['std']:0.3f}s"
+            )
+        self.log(
+            f"{' ' * 4}CPU mem: {stats['cpu_mem_use']['avg'] / (1024**2):0.3f} MB, GPU mem: {stats['gpu_mem_use']['avg'] / (1024**2):0.3f} MB"
+        )
+
+    def get_current_process_gpu_memory(self):
+        # Get all processes running on the GPU
+        processes = pynvml.nvmlDeviceGetComputeRunningProcesses(self.handle)
+
+        # Iterate through the processes to find the current process
+        for process in processes:
+            if process.pid == self.current_pid:
+                memory_usage = process.usedGpuMemory
+                return memory_usage

kaggle_environments/envs/lux_ai_s3/luxai_s3/pygame_render.py

@@ -0,0 +1,222 @@
+import numpy as np
+
+from luxai_s3.params import EnvParams
+from luxai_s3.state import ASTEROID_TILE, NEBULA_TILE, EnvState
+
+try:
+    import pygame
+except:
+    pass
+
+TILE_SIZE = 64
+
+
+class LuxAIPygameRenderer:
+    def __init__(self):
+        pass
+
+    def render(self, state: EnvState, params: EnvParams):
+        """Render the environment."""
+
+        # Initialize Pygame if not already done
+        if not pygame.get_init():
+            pygame.init()
+            self.clock = pygame.time.Clock()
+            # Set up the display
+            screen_width = params.map_width * TILE_SIZE
+            screen_height = params.map_height * TILE_SIZE
+            self.screen = pygame.display.set_mode((screen_width, screen_height))
+            self.surface = pygame.Surface(self.screen.get_size(), pygame.SRCALPHA)
+            pygame.display.set_caption("Lux AI Season 3")
+
+            self.display_options = {
+                "show_grid": True,
+                "show_relic_spots": False,
+                "show_sensor_mask": True,
+                "show_vision_power_map": True,
+                "show_energy_field": False,
+            }
+
+        # Handle events to keep the window responsive
+        render_state = "running"
+        while True:
+            self._update_display(state, params)
+            for event in pygame.event.get():
+                if event.type == pygame.TEXTINPUT:
+                    if event.text == " ":
+                        if render_state == "running":
+                            render_state = "paused"
+                        else:
+                            render_state = "running"
+                    elif event.text == "r":
+                        self.display_options["show_relic_spots"] = not self.display_options["show_relic_spots"]
+                    elif event.text == "s":
+                        self.display_options["show_sensor_mask"] = not self.display_options["show_sensor_mask"]
+                    elif event.text == "e":
+                        self.display_options["show_energy_field"] = not self.display_options["show_energy_field"]
+            else:
+                if render_state == "paused":
+                    self.clock.tick(60)
+                    continue
+                break
+
+    def _update_display(self, state: EnvState, params: EnvParams):
+        # Fill the screen with a background color
+        self.screen.fill((200, 200, 200))
+        self.surface.fill((200, 200, 200, 255))  # Light gray background
+
+        # Draw the grid of tiles
+        for x in range(params.map_width):
+            for y in range(params.map_height):
+                rect = pygame.Rect(x * TILE_SIZE, y * TILE_SIZE, TILE_SIZE, TILE_SIZE)
+                tile_type = state.map_features.tile_type[x, y]
+                if tile_type == NEBULA_TILE:
+                    color = (166, 177, 225, 255)  # Light blue (a6b1e1) for tile type 1
+                elif tile_type == ASTEROID_TILE:
+                    color = (51, 56, 68, 255)
+                else:
+                    color = (255, 255, 255, 255)  # White for other tile types
+                pygame.draw.rect(self.surface, color, rect)  # Draw filled squares
+
+        # Draw relic node configs if display option is enabled
+        def draw_rect_alpha(surface, color, rect):
+            shape_surf = pygame.Surface(pygame.Rect(rect).size, pygame.SRCALPHA)
+            pygame.draw.rect(shape_surf, color, shape_surf.get_rect())
+            surface.blit(shape_surf, rect)
+
+        if self.display_options["show_relic_spots"]:
+            mask = state.relic_nodes_map_weights
+            for x in range(params.map_width):
+                for y in range(params.map_height):
+                    if mask[x, y] > 0:
+                        rect = pygame.Rect(x * TILE_SIZE, y * TILE_SIZE, TILE_SIZE, TILE_SIZE)
+                        draw_rect_alpha(self.surface, (255, 215, 0, 50), rect)
+
+        # Draw energy nodes
+        for i in range(params.max_energy_nodes):
+            if state.energy_nodes_mask[i]:
+                x, y = state.energy_nodes[i, :2]
+                center_x = (x + 0.5) * TILE_SIZE
+                center_y = (y + 0.5) * TILE_SIZE
+                radius = TILE_SIZE // 4  # Adjust this value to change the size of the circle
+                pygame.draw.circle(
+                    self.surface,
+                    (0, 255, 0, 255),
+                    (int(center_x), int(center_y)),
+                    radius,
+                )
+        # Draw relic nodes
+        for i in range(params.max_relic_nodes):
+            if state.relic_nodes_mask[i]:
+                x, y = state.relic_nodes[i, :2]
+                rect_size = TILE_SIZE // 2  # Make the square smaller than the tile
+                rect_x = x * TILE_SIZE + (TILE_SIZE - rect_size) // 2
+                rect_y = y * TILE_SIZE + (TILE_SIZE - rect_size) // 2
+                rect = pygame.Rect(rect_x, rect_y, rect_size, rect_size)
+                pygame.draw.rect(self.surface, (173, 151, 32, 255), rect)  # Light blue color
+
+        # Draw sensor mask
+        if self.display_options["show_sensor_mask"]:
+            for team in range(params.num_teams):
+                for x in range(params.map_width):
+                    for y in range(params.map_height):
+                        if state.sensor_mask[team, x, y]:
+                            draw_rect_alpha(
+                                self.surface,
+                                (255, 0, 0, 25),
+                                pygame.Rect(x * TILE_SIZE, y * TILE_SIZE, TILE_SIZE, TILE_SIZE),
+                            )
+
+        if self.display_options["show_energy_field"]:
+            font = pygame.font.Font(None, 32)  # You may need to adjust the font size
+            for x in range(params.map_width):
+                for y in range(params.map_height):
+                    energy_field_value = state.map_features.energy[x, y]
+                    text = font.render(str(energy_field_value), True, (255, 255, 255))
+                    text_rect = text.get_rect(center=((x + 0.5) * TILE_SIZE, (y + 0.5) * TILE_SIZE))
+                    self.surface.blit(text, text_rect)
+                    if energy_field_value > 0:
+                        draw_rect_alpha(
+                            self.surface,
+                            (
+                                0,
+                                255,
+                                0,
+                                255 * energy_field_value / params.max_energy_per_tile,
+                            ),
+                            pygame.Rect(x * TILE_SIZE, y * TILE_SIZE, TILE_SIZE, TILE_SIZE),
+                        )
+                    else:
+                        draw_rect_alpha(
+                            self.surface,
+                            (
+                                255,
+                                0,
+                                0,
+                                255 * energy_field_value / params.min_energy_per_tile,
+                            ),
+                            pygame.Rect(x * TILE_SIZE, y * TILE_SIZE, TILE_SIZE, TILE_SIZE),
+                        )
+        # if self.display_options["show_vision_power_map"]:
+        #     print(state.vision_power_map.shape)
+        #     font = pygame.font.Font(None, 32)  # You may need to adjust the font size
+        #     # vision_power_map = vision_power_map - (state.map_features.tile_type == NEBULA_TILE)[..., 0] * params.nebula_tile_vision_reduction
+        #     for team in range(0, 1):
+        #         for x in range(params.map_width):
+        #             for y in range(params.map_height):
+        #                 vision_power_value = state.vision_power_map[team, x, y]
+        #                 vision_power_value -= state.map_features.tile_type[x, y] == NEBULA_TILE
+        #                 text = font.render(str(vision_power_value), True, (0, 255, 255))
+        #                 text_rect = text.get_rect(
+        #                     center=((x + 0.5) * TILE_SIZE, (y + 0.5) * TILE_SIZE)
+        #                 )
+        #                 self.surface.blit(text, text_rect)
+
+        # Draw units
+        for team in range(2):
+            for i in range(params.max_units):
+                if state.units_mask[team, i]:
+                    x, y = state.units.position[team, i]
+                    center_x = (x + 0.5) * TILE_SIZE
+                    center_y = (y + 0.5) * TILE_SIZE
+                    radius = TILE_SIZE // 3  # Adjust this value to change the size of the circle
+                    color = (255, 0, 0, 255) if team == 0 else (0, 0, 255, 255)  # Red for team 0, Blue for team 1
+                    pygame.draw.circle(self.surface, color, (int(center_x), int(center_y)), radius)
+        # Draw unit counts
+        unit_counts = {}
+        for team in range(2):
+            for i in range(params.max_units):
+                if state.units_mask[team, i]:
+                    x, y = np.array(state.units.position[team, i])
+                    pos = (x, y)
+                    if pos not in unit_counts:
+                        unit_counts[pos] = 0
+                    unit_counts[pos] += 1
+
+        font = pygame.font.Font(None, 32)  # You may need to adjust the font size
+        for pos, count in unit_counts.items():
+            if count >= 1:
+                x, y = pos
+                text = font.render(str(count), True, (255, 255, 255))  # White text
+                text_rect = text.get_rect(center=((x + 0.5) * TILE_SIZE, (y + 0.5) * TILE_SIZE))
+                self.surface.blit(text, text_rect)
+
+        # Draw the grid lines
+        for x in range(params.map_width + 1):
+            pygame.draw.line(
+                self.surface,
+                (100, 100, 100),
+                (x * TILE_SIZE, 0),
+                (x * TILE_SIZE, params.map_height * TILE_SIZE),
+            )
+        for y in range(params.map_height + 1):
+            pygame.draw.line(
+                self.surface,
+                (100, 100, 100),
+                (0, y * TILE_SIZE),
+                (params.map_width * TILE_SIZE, y * TILE_SIZE),
+            )
+
+        self.screen.blit(self.surface, (0, 0))
+        # Update the display
+        pygame.display.flip()

kaggle_environments/envs/lux_ai_s3/luxai_s3/spaces.py

@@ -0,0 +1,27 @@
+import chex
+import jax
+import jax.numpy as jnp
+import numpy as np
+from gymnax.environments.spaces import Space
+
+
+class MultiDiscrete(Space):
+    """Minimal jittable class for multi discrete gymnax spaces."""
+
+    def __init__(self, low: np.ndarray, high: np.ndarray):
+        self.low = low
+        self.high = high
+        self.dist = self.high - self.low
+        assert low.shape == high.shape
+        self.shape = low.shape
+        self.dtype = jnp.int16
+
+    def sample(self, rng: chex.PRNGKey) -> chex.Array:
+        return (jax.random.uniform(rng, shape=self.shape, minval=0, maxval=1) * self.dist + self.low).astype(self.dtype)
+
+    def contains(self, x) -> jnp.ndarray:
+        """Check whether specific object is within space."""
+        # type_cond = isinstance(x, self.dtype)
+        # shape_cond = (x.shape == self.shape)
+        range_cond = jnp.logical_and(x >= 0, x < self.n)
+        return range_cond