kaggle-environments 1.15.2__py2.py3-none-any.whl → 1.16.0__py2.py3-none-any.whl

kaggle_environments/envs/lux_ai_s3/luxai_s3/state.py

@@ -0,0 +1,399 @@
+import functools
+import chex
+import flax
+import jax
+import jax.numpy as jnp
+import numpy as np
+from flax import struct
+
+from luxai_s3.params import MAP_TYPES, EnvParams
+from luxai_s3.utils import to_numpy
+EMPTY_TILE = 0
+NEBULA_TILE = 1
+ASTEROID_TILE = 2
+
+ENERGY_NODE_FNS = [
+    lambda d, x, y, z: jnp.sin(d * x + y) * z, lambda d, x, y, z: (x / (d + 1) + y) * z
+]
+
+@struct.dataclass
+class UnitState:
+    position: chex.Array
+    """Position of the unit with shape (2) for x, y"""
+    energy: int
+    """Energy of the unit"""
+
+@struct.dataclass
+class MapTile:
+    energy: int
+    """Energy of the tile, generated via energy_nodes and energy_node_fns"""
+    tile_type: int
+    """Type of the tile"""
+
+@struct.dataclass
+class EnvState:
+    units: UnitState
+    """Units in the environment with shape (T, N, 3) for T teams, N max units, and 3 features.
+
+    3 features are for position (x, y), and energy
+    """
+    units_mask: chex.Array
+    """Mask of units in the environment with shape (T, N) for T teams, N max units"""
+    energy_nodes: chex.Array
+    """Energy nodes in the environment with shape (N, 2) for N max energy nodes, and 2 features.
+
+    2 features are for position (x, y)
+    """
+    
+    energy_node_fns: chex.Array
+    """Energy node functions for computing the energy field of the map. They describe the function with a sequence of numbers
+    
+    The first number is the function used. The subsequent numbers parameterize the function. The function is applied to distance of map tile to energy node and the function parameters.
+    """
+
+    # energy_field: chex.Array
+    # """Energy field in the environment with shape (H, W) for H height, W width. This is generated from other state"""
+    
+    energy_nodes_mask: chex.Array
+    """Mask of energy nodes in the environment with shape (N) for N max energy nodes"""
+    relic_nodes: chex.Array
+    """Relic nodes in the environment with shape (N, 2) for N max relic nodes, and 2 features.
+
+    2 features are for position (x, y)
+    """
+    relic_node_configs: chex.Array
+    """Relic node configs in the environment with shape (N, K, K) for N max relic nodes and a KxK relic configuration"""
+    relic_nodes_mask: chex.Array
+    """Mask of relic nodes in the environment with shape (N, ) for N max relic nodes"""
+    relic_nodes_map_weights: chex.Array
+    """Map of relic nodes in the environment with shape (H, W) for H height, W width. True if a relic node is present, False otherwise. This is generated from other state"""
+
+    map_features: MapTile
+    """Map features in the environment with shape (W, H, 2) for W width, H height
+    """
+
+    sensor_mask: chex.Array
+    """Sensor mask in the environment with shape (T, H, W) for T teams, H height, W width. This is generated from other state"""
+
+    vision_power_map: chex.Array
+    """Vision power map in the environment with shape (T, H, W) for T teams, H height, W width. This is generated from other state"""
+
+    team_points: chex.Array
+    """Team points in the environment with shape (T) for T teams"""
+    team_wins: chex.Array
+    """Team wins in the environment with shape (T) for T teams"""
+
+    steps: int = 0
+    """steps taken in the environment"""
+    match_steps: int = 0
+    """steps taken in the current match"""
+    
+@struct.dataclass
+class EnvObs:
+    """Partial observation of environment"""
+    units: UnitState
+    """Units in the environment with shape (T, N, 3) for T teams, N max units, and 3 features.
+
+    3 features are for position (x, y), and energy
+    """
+    units_mask: chex.Array
+    """Mask of units in the environment with shape (T, N) for T teams, N max units"""
+    
+    sensor_mask: chex.Array
+    
+    map_features: MapTile
+    """Map features in the environment with shape (W, H, 2) for W width, H height
+    """
+    relic_nodes: chex.Array
+    """Position of all relic nodes with shape (N, 2) for N max relic nodes and 2 features for position (x, y). Number is -1 if not visible"""
+    relic_nodes_mask: chex.Array
+    """Mask of all relic nodes with shape (N) for N max relic nodes"""
+    team_points: chex.Array
+    """Team points in the environment with shape (T) for T teams"""
+    team_wins: chex.Array
+    """Team wins in the environment with shape (T) for T teams"""
+    steps: int = 0
+    """steps taken in the environment"""
+    match_steps: int = 0
+    """steps taken in the current match"""
+    
+    
+
+def serialize_env_states(env_states: list[EnvState]):
+    def serialize_array(root: EnvState, arr, key_path: str = ""):
+        if key_path in ["sensor_mask", "relic_nodes_mask", "energy_nodes_mask", "energy_node_fns", "relic_nodes_map_weights"]:
+            return None
+        if key_path == "relic_nodes":
+            return root.relic_nodes[root.relic_nodes_mask].tolist()
+        if key_path == "relic_node_configs":
+            return root.relic_node_configs[root.relic_nodes_mask].tolist()
+        if key_path == "energy_nodes":
+            return root.energy_nodes[root.energy_nodes_mask].tolist()
+        if isinstance(arr, jnp.ndarray):
+            return arr.tolist()
+        elif isinstance(arr, dict):
+            ret = dict()
+            for k, v in arr.items():
+                new_key = key_path + "/" + k if key_path else k
+                new_val = serialize_array(root, v, new_key)
+                if new_val is not None:
+                    ret[k] = new_val
+            return ret
+        return arr
+    steps = []
+    for state in env_states:
+        state_dict = flax.serialization.to_state_dict(state)
+        steps.append(serialize_array(state, state_dict))
+
+    return steps
+
+def serialize_env_actions(env_actions: list):
+    def serialize_array(arr, key_path: str = ""):
+        if isinstance(arr, np.ndarray):
+            return arr.tolist()
+        elif isinstance(arr, jnp.ndarray):
+            return arr.tolist()
+        elif isinstance(arr, dict):
+            ret = dict()
+            for k, v in arr.items():
+                new_key = key_path + "/" + k if key_path else k
+                new_val = serialize_array(v, new_key)
+                if new_val is not None:
+                    ret[k] = new_val
+            return ret
+
+        return arr
+    steps = []
+    for state in env_actions:
+        state = flax.serialization.to_state_dict(state)
+        steps.append(serialize_array(state))
+
+    return steps
+
+def state_to_flat_obs(state: EnvState) -> chex.Array:
+    pass
+
+
+def flat_obs_to_state(flat_obs: chex.Array) -> EnvState:
+    pass
+
+@functools.partial(jax.jit, static_argnums=(2, 3, 4, 5, 6, 7, 8, 9))
+def gen_state(key: chex.PRNGKey, env_params: EnvParams, max_units: int, num_teams: int, map_type: int, map_width: int, map_height: int, max_energy_nodes: int, max_relic_nodes: int, relic_config_size: int) -> EnvState:
+    generated = gen_map(key, env_params, map_type, map_width, map_height, max_energy_nodes, max_relic_nodes, relic_config_size)
+    relic_nodes_map_weights = jnp.zeros(
+        shape=(map_width, map_height), dtype=jnp.int16
+    )
+
+    # TODO (this could be optimized better)
+    def update_relic_node(relic_nodes_map_weights, relic_data):
+        relic_node, relic_node_config, mask = relic_data
+        start_y = relic_node[1] - relic_config_size // 2
+        start_x = relic_node[0] - relic_config_size // 2
+        for dy in range(relic_config_size):
+            for dx in range(relic_config_size):
+                y, x = start_y + dy, start_x + dx
+                valid_pos = jnp.logical_and(
+                    jnp.logical_and(y >= 0, x >= 0),
+                    jnp.logical_and(y < map_height, x < map_width),
+                )
+                relic_nodes_map_weights = jnp.where(
+                    valid_pos & mask,
+                    relic_nodes_map_weights.at[x, y].add(relic_node_config[dx, dy].astype(jnp.int16)),
+                    relic_nodes_map_weights,
+                )
+        return relic_nodes_map_weights, None
+
+    # this is really slow...
+    relic_nodes_map_weights, _ = jax.lax.scan(
+        update_relic_node,
+        relic_nodes_map_weights,
+        (
+            generated["relic_nodes"],
+            generated["relic_node_configs"],
+            generated["relic_nodes_mask"],
+        ),
+    )
+    state = EnvState(
+        units=UnitState(position=jnp.zeros(shape=(num_teams, max_units, 2), dtype=jnp.int16), energy=jnp.zeros(shape=(num_teams, max_units, 1), dtype=jnp.int16)),
+        units_mask=jnp.zeros(
+            shape=(num_teams, max_units), dtype=jnp.bool
+        ),
+        team_points=jnp.zeros(shape=(num_teams), dtype=jnp.int32),
+        team_wins=jnp.zeros(shape=(num_teams), dtype=jnp.int32),
+        energy_nodes=generated["energy_nodes"],
+        energy_node_fns=generated["energy_node_fns"],
+        energy_nodes_mask=generated["energy_nodes_mask"],
+        # energy_field=jnp.zeros(shape=(params.map_height, params.map_width), dtype=jnp.int16),
+        relic_nodes=generated["relic_nodes"],
+        relic_nodes_mask=generated["relic_nodes_mask"],
+        relic_node_configs=generated["relic_node_configs"],
+        relic_nodes_map_weights=relic_nodes_map_weights,
+        sensor_mask=jnp.zeros(
+            shape=(num_teams, map_height, map_width),
+            dtype=jnp.bool,
+        ),
+        vision_power_map=jnp.zeros(shape=(num_teams, map_height, map_width), dtype=jnp.int16),
+        map_features=generated["map_features"],
+    )
+    return state
+
+@functools.partial(jax.jit, static_argnums=(2, 3, 4, 5, 6, 7))
+def gen_map(key: chex.PRNGKey, params: EnvParams, map_type: int, map_height: int, map_width: int, max_energy_nodes: int, max_relic_nodes: int, relic_config_size: int) -> chex.Array:
+    map_features = MapTile(energy=jnp.zeros(
+        shape=(map_height, map_width), dtype=jnp.int16
+    ), tile_type=jnp.zeros(
+        shape=(map_height, map_width), dtype=jnp.int16
+    ))
+    energy_nodes = jnp.zeros(shape=(max_energy_nodes, 2), dtype=jnp.int16)
+    energy_nodes_mask = jnp.zeros(shape=(max_energy_nodes), dtype=jnp.bool)
+    relic_nodes = jnp.zeros(shape=(max_relic_nodes, 2), dtype=jnp.int16)
+    relic_nodes_mask = jnp.zeros(shape=(max_relic_nodes), dtype=jnp.bool)
+    
+    if MAP_TYPES[map_type] == "random":
+        
+        ### Generate nebula tiles ###
+        key, subkey = jax.random.split(key)
+        perlin_noise = generate_perlin_noise_2d(subkey, (map_height, map_width), (4, 4))
+        noise = jnp.where(perlin_noise > 0.5, 1, 0)
+        # mirror along diagonal
+        noise = noise | noise.T
+        noise = noise[::-1, ::1]
+        map_features = map_features.replace(tile_type=jnp.where(noise, NEBULA_TILE, 0))
+        
+        ### Generate asteroid tiles ###
+        key, subkey = jax.random.split(key)
+        perlin_noise = generate_perlin_noise_2d(subkey, (map_height, map_width), (8, 8))
+        noise = jnp.where(perlin_noise < -0.5, 1, 0)
+        # mirror along diagonal
+        noise = noise | noise.T
+        noise = noise[::-1, ::1]
+        map_features = map_features.replace(tile_type=jnp.place(map_features.tile_type, noise, ASTEROID_TILE, inplace=False))
+        
+        ### Generate relic nodes ###
+        key, subkey = jax.random.split(key)
+        noise = generate_perlin_noise_2d(subkey, (map_height, map_width), (4, 4))
+        # Find the positions of the  highest noise values
+        flat_indices = jnp.argsort(noise.ravel())[-max_relic_nodes // 2:]  # Get indices of two highest values
+        highest_positions = jnp.column_stack(jnp.unravel_index(flat_indices, noise.shape))
+
+        # relic nodes have a fixed density of 25% nearby tiles can yield points
+        relic_node_configs = (
+            jax.random.randint(
+                key,
+                shape=(
+                    max_relic_nodes,
+                    relic_config_size,
+                    relic_config_size,
+                ),
+                minval=0,
+                maxval=10,
+            ).astype(jnp.float32)
+            >= 7.5
+        )
+        highest_positions = highest_positions.astype(jnp.int16)
+        relic_nodes_mask = relic_nodes_mask.at[0].set(True)
+        relic_nodes_mask = relic_nodes_mask.at[1].set(True)
+        mirrored_positions = jnp.stack([map_width - highest_positions[:, 1] - 1, map_height - highest_positions[:, 0] - 1], dtype=jnp.int16, axis=-1)
+        relic_nodes = jnp.concat([highest_positions, mirrored_positions], axis=0)
+        
+        key, subkey = jax.random.split(key)
+        relic_nodes_mask_half = jax.random.randint(key, (max_relic_nodes // 2, ), minval=0, maxval=2).astype(jnp.bool)
+        relic_nodes_mask_half = relic_nodes_mask_half.at[0].set(True)
+        relic_nodes_mask = relic_nodes_mask.at[:max_relic_nodes // 2].set(relic_nodes_mask_half)
+        relic_nodes_mask = relic_nodes_mask.at[max_relic_nodes // 2:].set(relic_nodes_mask_half)
+        # import ipdb;ipdb.set_trace()
+        relic_node_configs = relic_node_configs.at[max_relic_nodes // 2:].set(relic_node_configs[:max_relic_nodes // 2].transpose(0, 2, 1)[:, ::-1, ::-1])
+        
+        ### Generate energy nodes ###
+        key, subkey = jax.random.split(key)
+        noise = generate_perlin_noise_2d(subkey, (map_height, map_width), (4, 4))
+        # Find the positions of the  highest noise values
+        flat_indices = jnp.argsort(noise.ravel())[-max_energy_nodes // 2:]  # Get indices of highest values
+        highest_positions = jnp.column_stack(jnp.unravel_index(flat_indices, noise.shape)).astype(jnp.int16)
+        mirrored_positions = jnp.stack([map_width - highest_positions[:, 1] - 1, map_height - highest_positions[:, 0] - 1], dtype=jnp.int16, axis=-1)
+        energy_nodes = jnp.concat([highest_positions, mirrored_positions], axis=0)
+        key, subkey = jax.random.split(key)
+        energy_nodes_mask_half = jax.random.randint(key, (max_energy_nodes // 2, ), minval=0, maxval=2).astype(jnp.bool)
+        energy_nodes_mask_half = energy_nodes_mask_half.at[0].set(True)
+        energy_nodes_mask = energy_nodes_mask.at[:max_energy_nodes // 2].set(energy_nodes_mask_half)
+        energy_nodes_mask = energy_nodes_mask.at[max_energy_nodes // 2:].set(energy_nodes_mask_half)
+        
+        # TODO (stao): provide more randomization options for energy node functions.
+        energy_node_fns = jnp.array(
+            [
+                [0, 1.2, 1, 4],
+                [0, 0, 0, 0],
+                [0, 0, 0, 0],
+                # [1, 4, 0, 2],
+                [0, 1.2, 1, 4],
+                [0, 0, 0, 0],
+                [0, 0, 0, 0],
+                # [1, 4, 0, 0]
+            ]
+        )
+        # import ipdb; ipdb.set_trace()
+        # energy_node_fns = jnp.concat([energy_node_fns, jnp.zeros((params.max_energy_nodes - 2, 4), dtype=jnp.float32)], axis=0)
+        
+        
+    return dict(
+        map_features=map_features,
+        energy_nodes=energy_nodes,
+        energy_node_fns=energy_node_fns,
+        relic_nodes=relic_nodes,
+        energy_nodes_mask=energy_nodes_mask,
+        relic_nodes_mask=relic_nodes_mask,
+        relic_node_configs=relic_node_configs,
+    )
+def interpolant(t):
+    return t*t*t*(t*(t*6 - 15) + 10)
+
+@functools.partial(jax.jit, static_argnums=(1, 2, 3, 4))
+def generate_perlin_noise_2d(
+    key, shape, res, tileable=(False, False), interpolant=interpolant
+):
+    """Generate a 2D numpy array of perlin noise.
+
+    Args:
+        shape: The shape of the generated array (tuple of two ints).
+            This must be a multple of res.
+        res: The number of periods of noise to generate along each
+            axis (tuple of two ints). Note shape must be a multiple of
+            res.
+        tileable: If the noise should be tileable along each axis
+            (tuple of two bools). Defaults to (False, False).
+        interpolant: The interpolation function, defaults to
+            t*t*t*(t*(t*6 - 15) + 10).
+
+    Returns:
+        A numpy array of shape shape with the generated noise.
+
+    Raises:
+        ValueError: If shape is not a multiple of res.
+    """
+    delta = (res[0] / shape[0], res[1] / shape[1])
+    d = (shape[0] // res[0], shape[1] // res[1])
+    grid = jnp.mgrid[0:res[0]:delta[0], 0:res[1]:delta[1]]\
+             .transpose(1, 2, 0) % 1
+    # Gradients
+    angles = 2*jnp.pi*jax.random.uniform(key, (res[0]+1, res[1]+1))
+    gradients = jnp.dstack((jnp.cos(angles), jnp.sin(angles)))
+    if tileable[0]:
+        gradients[-1,:] = gradients[0,:]
+    if tileable[1]:
+        gradients[:,-1] = gradients[:,0]
+    gradients = gradients.repeat(d[0], 0).repeat(d[1], 1)
+    g00 = gradients[    :-d[0],    :-d[1]]
+    g10 = gradients[d[0]:     ,    :-d[1]]
+    g01 = gradients[    :-d[0],d[1]:     ]
+    g11 = gradients[d[0]:     ,d[1]:     ]
+    
+    # Ramps
+    n00 = jnp.sum(jnp.dstack((grid[:,:,0]  , grid[:,:,1]  )) * g00, 2)
+    n10 = jnp.sum(jnp.dstack((grid[:,:,0]-1, grid[:,:,1]  )) * g10, 2)
+    n01 = jnp.sum(jnp.dstack((grid[:,:,0]  , grid[:,:,1]-1)) * g01, 2)
+    n11 = jnp.sum(jnp.dstack((grid[:,:,0]-1, grid[:,:,1]-1)) * g11, 2)
+    # Interpolation
+    t = interpolant(grid)
+    n0 = n00*(1-t[:,:,0]) + t[:,:,0]*n10
+    n1 = n01*(1-t[:,:,0]) + t[:,:,0]*n11
+    return jnp.sqrt(2)*((1-t[:,:,1])*n0 + t[:,:,1]*n1)

kaggle_environments/envs/lux_ai_s3/luxai_s3/utils.py

@@ -0,0 +1,12 @@
+import numpy as np
+
+
+def to_numpy(x):
+    if isinstance(x, dict):
+        return {k: to_numpy(v) for k, v in x.items()}
+    elif isinstance(x, list):
+        return np.array(x)
+    elif isinstance(x, np.ndarray):
+        return x
+    else:
+        return np.array(x)

kaggle_environments/envs/lux_ai_s3/luxai_s3/wrappers.py

@@ -0,0 +1,187 @@
+# TODO (stao): Add lux ai s3 env to gymnax api wrapper, which is the old gym api
+import json
+import os
+from typing import Any, SupportsFloat
+import flax
+import flax.serialization
+import gymnasium as gym
+import gymnax
+import gymnax.environments.spaces
+import jax
+import numpy as np
+import dataclasses
+from luxai_s3.env import LuxAIS3Env
+from luxai_s3.params import EnvParams, env_params_ranges
+from luxai_s3.state import serialize_env_actions, serialize_env_states
+from luxai_s3.utils import to_numpy
+
+
+class LuxAIS3GymEnv(gym.Env):
+    def __init__(self, numpy_output: bool = False):
+        self.numpy_output = numpy_output
+        self.rng_key = jax.random.key(0)
+        self.jax_env = LuxAIS3Env(auto_reset=False)
+        self.env_params: EnvParams = EnvParams()
+
+        # auto run compiling steps here:
+        # print("Running compilation steps")
+        key = jax.random.key(0)
+        # Reset the environment
+        dummy_env_params = EnvParams(map_type=1)
+        key, reset_key = jax.random.split(key)
+        obs, state = self.jax_env.reset(reset_key, params=dummy_env_params)
+        # Take a random action
+        key, subkey = jax.random.split(key)
+        action = self.jax_env.action_space(dummy_env_params).sample(subkey)
+        # Step the environment and compile. Not sure why 2 steps? are needed
+        for _ in range(2):
+            key, subkey = jax.random.split(key)
+            obs, state, reward, terminated, truncated, info = self.jax_env.step(
+                subkey, state, action, params=dummy_env_params
+            )
+        # print("Finish compilation steps")
+        low = np.zeros((self.env_params.max_units, 3))
+        low[:, 1:] = -self.env_params.unit_sap_range
+        high = np.ones((self.env_params.max_units, 3)) * 6
+        high[:, 1:] = self.env_params.unit_sap_range
+        self.action_space = gym.spaces.Dict(
+            dict(
+                player_0=gym.spaces.Box(low=low, high=high, dtype=np.int16),
+                player_1=gym.spaces.Box(low=low, high=high, dtype=np.int16),
+            )
+        )
+
+    def render(self):
+        self.jax_env.render(self.state, self.env_params)
+
+    def reset(
+        self, *, seed: int | None = None, options: dict[str, Any] | None = None
+    ) -> tuple[Any, dict[str, Any]]:
+        if seed is not None:
+            self.rng_key = jax.random.key(seed)
+        self.rng_key, reset_key = jax.random.split(self.rng_key)
+        # generate random game parameters
+        # TODO (stao): check why this keeps recompiling when marking structs as static args
+        randomized_game_params = dict()
+        for k, v in env_params_ranges.items():
+            self.rng_key, subkey = jax.random.split(self.rng_key)
+            randomized_game_params[k] = jax.random.choice(
+                subkey, jax.numpy.array(v)
+            ).item()
+        params = EnvParams(**randomized_game_params)
+        if options is not None and "params" in options:
+            params = options["params"]
+
+        self.env_params = params
+        obs, self.state = self.jax_env.reset(reset_key, params=params)
+        if self.numpy_output:
+            obs = to_numpy(flax.serialization.to_state_dict(obs))
+
+        # only keep the following game parameters available to the agent
+        params_dict = dataclasses.asdict(params)
+        params_dict_kept = dict()
+        for k in [
+            "max_units",
+            "match_count_per_episode",
+            "max_steps_in_match",
+            "map_height",
+            "map_width",
+            "num_teams",
+            "unit_move_cost",
+            "unit_sap_cost",
+            "unit_sap_range",
+            "unit_sensor_range",
+        ]:
+            params_dict_kept[k] = params_dict[k]
+        return obs, dict(
+            params=params_dict_kept, full_params=params_dict, state=self.state
+        )
+
+    def step(
+        self, action: Any
+    ) -> tuple[Any, SupportsFloat, bool, bool, dict[str, Any]]:
+        self.rng_key, step_key = jax.random.split(self.rng_key)
+        obs, self.state, reward, terminated, truncated, info = self.jax_env.step(
+            step_key, self.state, action, self.env_params
+        )
+        if self.numpy_output:
+            obs = to_numpy(flax.serialization.to_state_dict(obs))
+            reward = to_numpy(reward)
+            terminated = to_numpy(terminated)
+            truncated = to_numpy(truncated)
+            # info = to_numpy(flax.serialization.to_state_dict(info))
+        return obs, reward, terminated, truncated, info
+
+
+# TODO: vectorized gym wrapper
+
+
+class RecordEpisode(gym.Wrapper):
+    def __init__(
+        self,
+        env: LuxAIS3GymEnv,
+        save_dir: str = None,
+        save_on_close: bool = True,
+        save_on_reset: bool = True,
+    ):
+        super().__init__(env)
+        self.episode = dict(states=[], actions=[], metadata=dict())
+        self.episode_id = 0
+        self.save_dir = save_dir
+        self.save_on_close = save_on_close
+        self.save_on_reset = save_on_reset
+        self.episode_steps = 0
+        if save_dir is not None:
+            from pathlib import Path
+
+            Path(save_dir).mkdir(parents=True, exist_ok=True)
+
+    def reset(
+        self, *, seed: int | None = None, options: dict[str, Any] | None = None
+    ) -> tuple[Any, dict[str, Any]]:
+        if self.save_on_reset and self.episode_steps > 0:
+            self._save_episode_and_reset()
+        obs, info = self.env.reset(seed=seed, options=options)
+
+        self.episode["metadata"]["seed"] = seed
+        self.episode["params"] = flax.serialization.to_state_dict(info["full_params"])
+        self.episode["states"].append(info["state"])
+        return obs, info
+
+    def step(
+        self, action: Any
+    ) -> tuple[Any, SupportsFloat, bool, bool, dict[str, Any]]:
+        obs, reward, terminated, truncated, info = self.env.step(action)
+        self.episode_steps += 1
+        self.episode["states"].append(info["final_state"])
+        self.episode["actions"].append(action)
+        return obs, reward, terminated, truncated, info
+
+    def serialize_episode_data(self, episode=None):
+        if episode is None:
+            episode = self.episode
+        ret = dict()
+        ret["observations"] = serialize_env_states(episode["states"])
+        if "actions" in episode:
+            ret["actions"] = serialize_env_actions(episode["actions"])
+        ret["metadata"] = episode["metadata"]
+        ret["params"] = episode["params"]
+        return ret
+
+    def save_episode(self, save_path: str):
+        episode = self.serialize_episode_data()
+        with open(save_path, "w") as f:
+            json.dump(episode, f)
+        self.episode = dict(states=[], actions=[], metadata=dict())
+
+    def _save_episode_and_reset(self):
+        """saves to generated path based on self.save_dir and episoe id and updates relevant counters"""
+        self.save_episode(
+            os.path.join(self.save_dir, f"episode_{self.episode_id}.json")
+        )
+        self.episode_id += 1
+        self.episode_steps = 0
+
+    def close(self):
+        if self.save_on_close and self.episode_steps > 0:
+            self._save_episode_and_reset()

kaggle_environments/envs/lux_ai_s3/test_agents/python/agent.py

@@ -0,0 +1,71 @@
+if __package__ == "":
+    from lux.utils import direction_to
+else:
+    from .lux.utils import direction_to
+import numpy as np
+class Agent():
+    def __init__(self, player: str, env_cfg) -> None:
+        self.player = player
+        self.opp_player = "player_1" if self.player == "player_0" else "player_0"
+        self.team_id = 0 if self.player == "player_0" else 1
+        self.opp_team_id = 1 if self.team_id == 0 else 0
+        np.random.seed(0)
+        self.env_cfg = env_cfg
+        
+        self.relic_node_positions = []
+        self.discovered_relic_nodes_ids = set()
+        self.unit_explore_locations = dict()
+
+    def act(self, step: int, obs, remainingOverageTime: int = 60):
+        """implement this function to decide what actions to send to each available unit. 
+        
+        step is the current timestep number of the game starting from 0 going up to max_steps_in_match * match_count_per_episode - 1.
+        """
+        unit_mask = np.array(obs["units_mask"][self.team_id]) # shape (max_units, )
+        unit_positions = np.array(obs["units"]["position"][self.team_id]) # shape (max_units, 2)
+        unit_energys = np.array(obs["units"]["energy"][self.team_id]) # shape (max_units, 1)
+        observed_relic_node_positions = np.array(obs["relic_nodes"]) # shape (max_relic_nodes, 2)
+        observed_relic_nodes_mask = np.array(obs["relic_nodes_mask"]) # shape (max_relic_nodes, )
+        team_points = np.array(obs["team_points"]) # points of each team, team_points[self.team_id] is the points of the your team
+        
+        # ids of units you can control at this timestep
+        available_unit_ids = np.where(unit_mask)[0]
+        # visible relic nodes
+        visible_relic_node_ids = set(np.where(observed_relic_nodes_mask)[0])
+        
+        actions = np.zeros((self.env_cfg["max_units"], 3), dtype=int)
+
+
+        # basic strategy here is simply to have some units randomly explore and some units collecting as much energy as possible
+        # and once a relic node is found, we send all units to move randomly around the first relic node to gain points
+        # and information about where relic nodes are found are saved for the next match
+        
+        # save any new relic nodes that we discover for the rest of the game.
+        for id in visible_relic_node_ids:
+            if id not in self.discovered_relic_nodes_ids:
+                self.discovered_relic_nodes_ids.add(id)
+                self.relic_node_positions.append(observed_relic_node_positions[id])
+            
+
+        # unit ids range from 0 to max_units - 1
+        for unit_id in available_unit_ids:
+            unit_pos = unit_positions[unit_id]
+            unit_energy = unit_energys[unit_id]
+            if len(self.relic_node_positions) > 0:
+                nearest_relic_node_position = self.relic_node_positions[0]
+                manhattan_distance = abs(unit_pos[0] - nearest_relic_node_position[0]) + abs(unit_pos[1] - nearest_relic_node_position[1])
+                
+                # if close to the relic node we want to hover around it and hope to gain points
+                if manhattan_distance <= 4:
+                    random_direction = np.random.randint(0, 5)
+                    actions[unit_id] = [random_direction, 0, 0]
+                else:
+                    # otherwise we want to move towards the relic node
+                    actions[unit_id] = [direction_to(unit_pos, nearest_relic_node_position), 0, 0]
+            else:
+                # randomly explore by picking a random location on the map and moving there for about 20 steps
+                if step % 20 == 0 or unit_id not in self.unit_explore_locations:
+                    rand_loc = (np.random.randint(0, self.env_cfg["map_width"]), np.random.randint(0, self.env_cfg["map_height"]))
+                    self.unit_explore_locations[unit_id] = rand_loc
+                actions[unit_id] = [direction_to(unit_pos, self.unit_explore_locations[unit_id]), 0, 0]
+        return actions