Repository: Lux-AI-Challenge/Lux-Design-S2 (language: python)

Issue:
fix: actions correctly repeat now
The change log of v2.0.0 says, "actions in an action queue can specify both the number of times to repeat it directly, as well as whether to put it back to the end of the action queue once exhausted." However, it is not implemented in `unit.repeat_action`. This commit implements the such feature.


Discussion / hints:

Fix patch:
diff --git a/ChangeLog.md b/ChangeLog.md
index ac08aea1..8e88ca60 100644
--- a/ChangeLog.md
+++ b/ChangeLog.md
@@ -4,8 +4,23 @@
 
 Official release!
 
+Major Engine Changes
+- Weather is removed
+- Each tile of lichen owned and connected to a factory gives 1 power to the factory each turn
+- Actions in an action queue can specify both the number of times to repeat it directly, as well as whether to put it back to the end of the action queue once exhausted
+- When digging out lichen, if the final lichen is dug out then rubble equal to `DIG_RUBBLE_REMOVED` is added to the tile.
+- When handling collisions, if two units of the same weight class move onto the same tile, previously they both were destroyed. Now, all units with less power are destroyed and the unit with the most power in the collision loses half of the power of the unit with the 2nd most power.
+- Map generation has been updated to include resource distribution types, a matrix of low/high ore and low/high ice. Low = ~15 tiles, High = ~35 tiles.
+
+Configuration Changes
+- Light units rubble cost is floor(0.05 * rubble) now.
+- Heavy units dig cost is 60 power instead of 100.
+- It takes 20 lichen to grow to new tiles instead of 10 now.
+
+Repo Changes
 - Repository has been reorganized to accomodate more packages, specifically the gpu version
 
+
 ### v1.1.6
 - Fix bugs related to tournament tool and not printing some numbers as well as concurrent configuration not used
 ### v1.1.5
diff --git a/kits/README.md b/kits/README.md
index 4b7d092b..c49a8b7f 100644
--- a/kits/README.md
+++ b/kits/README.md
@@ -30,7 +30,7 @@ In the early phase, the action space is different than the actual game phase. Se
 
 During the actual game phase, factories have 3 possible actions, `build_light`, `build_heavy`, and `water`. Units/Robots (light or heavy robots) have 5 possible actions: `move`, `dig`, `transfer`, `pickup`, `self_destruct`, `recharge`; where `move, dig, self_destruct` have power costs
 
-In Lux AI Season 2, the robots's actual action space is a list of actions representing it's action queue and your agent will set this action queue to control robots. This action queue max size is `env_cfg.UNIT_ACTION_QUEUE_SIZE`. Each turn, the unit executes the action at the front of the queue. If the action is marked as to be repeated (the default), it is replaced to the back of the queue.
+In Lux AI Season 2, the robots's actual action space is a list of actions representing it's action queue and your agent will set this action queue to control robots. This action queue max size is `env_cfg.UNIT_ACTION_QUEUE_SIZE`. Each turn, the unit executes the action at the front of the queue, and repeatedly does this a user-specified `n` times until exhausted. If the action is marked as to be repeated, it is replaced to the back of the queue.
 
 In code, actions can be given to units as so
 
@@ -85,7 +85,6 @@ The general observation given to your bot in the kits will look like below. `Arr
       "valid_spawns_mask": Array(48, 48),
       "factories_per_team": int
     },
-    "weather": Array(1000),
     "real_env_steps": int,
     "teams": {
       "player_0": {
@@ -124,7 +123,6 @@ class GameState:
     env_steps: int # number of env steps passed
     env_cfg: dict # current env configuration
     board: Board # the game board
-    weather_schedule: Array # the weather ID at each time step
     units: Dict[str, Dict[str, Unit]] # maps agent ID (player_0, player_1) to a dictionary mapping unit ID to unit objects
     factories: Dict[str, Dict[str, Factory]] # maps agent ID (player_0, player_1) to a dictionary mapping unit ID to factory objects
     teams: Dict[str, Team] # maps agent ID (player_0, player_1) to a Team object
@@ -144,7 +142,7 @@ class Board:
     spawns: Array # possible spawn locations on your team's half of the map
 ```
 
-Each `Unit` object comes with functions to generate the action vector for actions like move and dig, as well as cost functions that return the power cost to perform some actions (and also factors weather into that cost)
+Each `Unit` object comes with functions to generate the action vector for actions like move and dig, as well as cost functions that return the power cost to perform some actions.
 
 Each `Factory` object comes with functions to generate actions as well as compute the cost of building robots and watering lichen.
 
diff --git a/kits/python/agent.py b/kits/python/agent.py
index 2bc0e4ea..e8f87f3f 100644
--- a/kits/python/agent.py
+++ b/kits/python/agent.py
@@ -43,9 +43,8 @@ def act(self, step: int, obs, remainingOverageTime: int = 60):
             if factory.power >= self.env_cfg.ROBOTS["HEAVY"].POWER_COST and \
             factory.cargo.metal >= self.env_cfg.ROBOTS["HEAVY"].METAL_COST:
                 actions[unit_id] = factory.build_heavy()
-            if self.env_cfg.max_episode_length - game_state.real_env_steps < 50:
-                if factory.water_cost(game_state) <= factory.cargo.water:
-                    actions[unit_id] = factory.water()
+            if factory.water_cost(game_state) <= factory.cargo.water / 5:
+                actions[unit_id] = factory.water()
             factory_tiles += [factory.pos]
             factory_units += [factory]
         factory_tiles = np.array(factory_tiles)
@@ -70,12 +69,12 @@ def act(self, step: int, obs, remainingOverageTime: int = 60):
                     closest_ice_tile = ice_tile_locations[np.argmin(ice_tile_distances)]
                     if np.all(closest_ice_tile == unit.pos):
                         if unit.power >= unit.dig_cost(game_state) + unit.action_queue_cost(game_state):
-                            actions[unit_id] = [unit.dig(repeat=0)]
+                            actions[unit_id] = [unit.dig(repeat=0, n=1)]
                     else:
                         direction = direction_to(unit.pos, closest_ice_tile)
                         move_cost = unit.move_cost(game_state, direction)
                         if move_cost is not None and unit.power >= move_cost + unit.action_queue_cost(game_state):
-                            actions[unit_id] = [unit.move(direction, repeat=0)]
+                            actions[unit_id] = [unit.move(direction, repeat=0, n=1)]
                 # else if we have enough ice, we go back to the factory and dump it.
                 elif unit.cargo.ice >= 40:
                     direction = direction_to(unit.pos, closest_factory_tile)
@@ -85,5 +84,5 @@ def act(self, step: int, obs, remainingOverageTime: int = 60):
                     else:
                         move_cost = unit.move_cost(game_state, direction)
                         if move_cost is not None and unit.power >= move_cost + unit.action_queue_cost(game_state):
-                            actions[unit_id] = [unit.move(direction, repeat=0)]
+                            actions[unit_id] = [unit.move(direction, repeat=0, n=1)]
         return actions
diff --git a/kits/python/lux/config.py b/kits/python/lux/config.py
index 97cf6b61..1ea9ef1b 100644
--- a/kits/python/lux/config.py
+++ b/kits/python/lux/config.py
@@ -1,6 +1,6 @@
+import dataclasses
 from argparse import Namespace
 from dataclasses import dataclass
-import dataclasses
 from typing import Dict, List
 
 
@@ -20,7 +20,7 @@ class UnitConfig:
     CHARGE: int = 1
     INIT_POWER: int = 50
     MOVE_COST: int = 1
-    RUBBLE_MOVEMENT_COST: int = 1
+    RUBBLE_MOVEMENT_COST: float = 1
     DIG_COST: int = 5
     DIG_RUBBLE_REMOVED: int = 1
     DIG_RESOURCE_GAIN: int = 2
@@ -30,17 +30,15 @@ class UnitConfig:
     ACTION_QUEUE_POWER_COST: int = 1
 
 
-
 @dataclass
 class EnvConfig:
     ## various options that can be configured if needed
 
-
     ### Variable parameters that don't affect game logic much ###
     max_episode_length: int = 1000
     map_size: int = 48
     verbose: int = 1
-    
+
     # this can be disabled to improve env FPS but assume your actions are well formatted
     # During online competition this is set to True
     validate_action_space: bool = True
@@ -52,19 +50,22 @@ class EnvConfig:
     MAX_FACTORIES: int = 5
     CYCLE_LENGTH: int = 50
     DAY_LENGTH: int = 30
-    UNIT_ACTION_QUEUE_SIZE: int = 20 # when set to 1, then no action queue is used
-
+    UNIT_ACTION_QUEUE_SIZE: int = 20  # when set to 1, then no action queue is used
 
     MAX_RUBBLE: int = 100
     FACTORY_RUBBLE_AFTER_DESTRUCTION: int = 50
-    INIT_WATER_METAL_PER_FACTORY: int = 150 # amount of water and metal units given to each factory
+    INIT_WATER_METAL_PER_FACTORY: int = (
+        150  # amount of water and metal units given to each factory
+    )
     INIT_POWER_PER_FACTORY: int = 1000
 
     #### LICHEN ####
-    MIN_LICHEN_TO_SPREAD: int = 10
+    MIN_LICHEN_TO_SPREAD: int = 20
     LICHEN_LOST_WITHOUT_WATER: int = 1
     LICHEN_GAINED_WITH_WATER: int = 1
     MAX_LICHEN_PER_TILE: int = 100
+    POWER_PER_CONNECTED_LICHEN_TILE: int = 1
+
     # cost of watering with a factory is `ceil(# of connected lichen tiles) / (this factor) + 1`
     LICHEN_WATERING_COST_FACTOR: int = 10
 
@@ -72,75 +73,64 @@ class EnvConfig:
     BIDDING_SYSTEM: bool = True
 
     #### Factories ####
-    FACTORY_PROCESSING_RATE_WATER: int = 50
-    ICE_WATER_RATIO: int = 2
+    FACTORY_PROCESSING_RATE_WATER: int = 100
+    ICE_WATER_RATIO: int = 4
     FACTORY_PROCESSING_RATE_METAL: int = 50
     ORE_METAL_RATIO: int = 5
     # game design note: Factories close to resource cluster = more resources are refined per turn
-    # Then the high ice:water and ore:metal ratios encourages transfer of refined resources between 
+    # Then the high ice:water and ore:metal ratios encourages transfer of refined resources between
     # factories dedicated to mining particular clusters which is more possible as it is more compact
-    # 
-
 
     FACTORY_CHARGE: int = 50
     FACTORY_WATER_CONSUMPTION: int = 1
     # game design note: with a positve water consumption, game becomes quite hard for new competitors.
     # so we set it to 0
 
+    #### Collision Mechanics ####
+    POWER_LOSS_FACTOR: float = 0.5
 
     #### Units ####
     ROBOTS: Dict[str, UnitConfig] = dataclasses.field(
         default_factory=lambda: dict(
-        LIGHT=UnitConfig(
-            METAL_COST=10, POWER_COST=50, INIT_POWER=50, CARGO_SPACE=100, BATTERY_CAPACITY=150, CHARGE=1, MOVE_COST=1, RUBBLE_MOVEMENT_COST=0,
-            DIG_COST=5,
-            SELF_DESTRUCT_COST=5,
-            DIG_RUBBLE_REMOVED=2,
-            DIG_RESOURCE_GAIN=2,
-            DIG_LICHEN_REMOVED=10,
-            RUBBLE_AFTER_DESTRUCTION=1,
-            ACTION_QUEUE_POWER_COST=1,
-        ),
-
-        HEAVY=UnitConfig(
-            METAL_COST=100, POWER_COST=500, INIT_POWER=500, CARGO_SPACE=1000, BATTERY_CAPACITY=3000, CHARGE=10, MOVE_COST=20, RUBBLE_MOVEMENT_COST=1,
-            DIG_COST=100,
-            SELF_DESTRUCT_COST=100,
-            DIG_RUBBLE_REMOVED=20,
-            DIG_RESOURCE_GAIN=20,
-            DIG_LICHEN_REMOVED=100,
-            RUBBLE_AFTER_DESTRUCTION=10,
-            ACTION_QUEUE_POWER_COST=10,
-        ),
+            LIGHT=UnitConfig(
+                METAL_COST=10,
+                POWER_COST=50,
+                INIT_POWER=50,
+                CARGO_SPACE=100,
+                BATTERY_CAPACITY=150,
+                CHARGE=1,
+                MOVE_COST=1,
+                RUBBLE_MOVEMENT_COST=0.05,
+                DIG_COST=5,
+                SELF_DESTRUCT_COST=5,
+                DIG_RUBBLE_REMOVED=2,
+                DIG_RESOURCE_GAIN=2,
+                DIG_LICHEN_REMOVED=10,
+                RUBBLE_AFTER_DESTRUCTION=1,
+                ACTION_QUEUE_POWER_COST=1,
+            ),
+            HEAVY=UnitConfig(
+                METAL_COST=100,
+                POWER_COST=500,
+                INIT_POWER=500,
+                CARGO_SPACE=1000,
+                BATTERY_CAPACITY=3000,
+                CHARGE=10,
+                MOVE_COST=20,
+                RUBBLE_MOVEMENT_COST=1,
+                DIG_COST=60,
+                SELF_DESTRUCT_COST=100,
+                DIG_RUBBLE_REMOVED=20,
+                DIG_RESOURCE_GAIN=20,
+                DIG_LICHEN_REMOVED=100,
+                RUBBLE_AFTER_DESTRUCTION=10,
+                ACTION_QUEUE_POWER_COST=10,
+            ),
         )
     )
 
-    #### Weather ####
-    WEATHER_ID_TO_NAME: list = dataclasses.field(default_factory=lambda : ["NONE","MARS_QUAKE","COLD_SNAP","DUST_STORM","SOLAR_FLARE"])
-    NUM_WEATHER_EVENTS_RANGE: List[int] = dataclasses.field(default_factory=lambda:[3,5])
-    WEATHER: dict = dataclasses.field(default_factory=lambda:dict(
-        MARS_QUAKE=dict(
-            TIME_RANGE=[1, 5]
-        ),
-        COLD_SNAP=dict(
-            # power multiplier required per robot action. 2 -> requires 2x as much power to execute the same action
-            POWER_CONSUMPTION=2,
-            TIME_RANGE=[10, 30]
-        ),
-        DUST_STORM=dict(
-            # power gain multiplier. .5 -> gain .5x as much power per turn
-            POWER_GAIN=0.5,
-            TIME_RANGE=[10, 30]
-        ),
-        SOLAR_FLARE=dict(
-            # power gain multiplier. 2 -> gain 2x as much power per turn
-            POWER_GAIN=2,
-            TIME_RANGE=[10, 30]
-        ),
-    ))
-
     @classmethod
     def from_dict(cls, data):
         data["ROBOTS"]["LIGHT"] = UnitConfig(**data["ROBOTS"]["LIGHT"])
         data["ROBOTS"]["HEAVY"] = UnitConfig(**data["ROBOTS"]["HEAVY"])
-        return cls(**data)
\ No newline at end of file
+        return cls(**data)
diff --git a/kits/python/lux/factory.py b/kits/python/lux/factory.py
index e66dda3f..13e87e82 100644
--- a/kits/python/lux/factory.py
+++ b/kits/python/lux/factory.py
@@ -2,7 +2,6 @@
 from sys import stderr
 import numpy as np
 from dataclasses import dataclass
-from lux.weather import get_weather_config
 from lux.cargo import UnitCargo
 from lux.config import EnvConfig
 @dataclass
@@ -21,10 +20,8 @@ def build_heavy_metal_cost(self, game_state):
         unit_cfg = self.env_cfg.ROBOTS["HEAVY"]
         return unit_cfg.METAL_COST
     def build_heavy_power_cost(self, game_state):
-        current_weather = game_state.weather_schedule[game_state.real_env_steps]
-        weather_cfg = get_weather_config(current_weather, self.env_cfg)
         unit_cfg = self.env_cfg.ROBOTS["HEAVY"]
-        return math.ceil(unit_cfg.POWER_COST * weather_cfg["power_loss_factor"])
+        return unit_cfg.POWER_COST
     def can_build_heavy(self, game_state):
         return self.power >= self.build_heavy_power_cost(game_state) and self.cargo.metal >= self.build_heavy_metal_cost(game_state)
     def build_heavy(self):
@@ -34,10 +31,8 @@ def build_light_metal_cost(self, game_state):
         unit_cfg = self.env_cfg.ROBOTS["LIGHT"]
         return unit_cfg.METAL_COST
     def build_light_power_cost(self, game_state):
-        current_weather = game_state.weather_schedule[game_state.real_env_steps]
-        weather_cfg = get_weather_config(current_weather, self.env_cfg)
         unit_cfg = self.env_cfg.ROBOTS["LIGHT"]
-        return math.ceil(unit_cfg.POWER_COST * weather_cfg["power_loss_factor"])
+        return unit_cfg.POWER_COST
     def can_build_light(self, game_state):
         return self.power >= self.build_light_power_cost(game_state) and self.cargo.metal >= self.build_light_metal_cost(game_state)
 
diff --git a/kits/python/lux/kit.py b/kits/python/lux/kit.py
index 0ab83024..8de70cd7 100644
--- a/kits/python/lux/kit.py
+++ b/kits/python/lux/kit.py
@@ -105,7 +105,6 @@ def obs_to_game_state(step, env_cfg: EnvConfig, obs):
             factories_per_team=obs["board"]["factories_per_team"],
             valid_spawns_mask=obs["board"]["valid_spawns_mask"]
         ),
-        weather_schedule=obs["weather_schedule"],
         units=units,
         factories=factories,
         teams=teams
@@ -130,7 +129,6 @@ class GameState:
     env_steps: int
     env_cfg: dict
     board: Board
-    weather_schedule: np.ndarray = None
     units: Dict[str, Dict[str, Unit]] = field(default_factory=dict)
     factories: Dict[str, Dict[str, Factory]] = field(default_factory=dict)
     teams: Dict[str, Team] = field(default_factory=dict)
diff --git a/kits/python/lux/unit.py b/kits/python/lux/unit.py
index 880ab505..8bba5ee1 100644
--- a/kits/python/lux/unit.py
+++ b/kits/python/lux/unit.py
@@ -3,7 +3,6 @@
 from typing import List
 import numpy as np
 from dataclasses import dataclass
-from lux.weather import get_weather_config
 from lux.cargo import UnitCargo
 from lux.config import EnvConfig
 
@@ -29,9 +28,7 @@ def agent_id(self):
 
     def action_queue_cost(self, game_state):
         cost = self.env_cfg.ROBOTS[self.unit_type].ACTION_QUEUE_POWER_COST
-        current_weather = game_state.weather_schedule[game_state.real_env_steps]
-        weather_cfg = get_weather_config(current_weather, self.env_cfg)
-        return cost * weather_cfg["power_loss_factor"]
+        return cost
 
     def move_cost(self, game_state, direction):
         board = game_state.board
@@ -45,41 +42,35 @@ def move_cost(self, game_state, direction):
             return None
         rubble_at_target = board.rubble[target_pos[0]][target_pos[1]]
         
-        current_weather = game_state.weather_schedule[game_state.real_env_steps]
-        weather_cfg = get_weather_config(current_weather, self.env_cfg)
-        return math.ceil((self.unit_cfg.MOVE_COST + self.unit_cfg.RUBBLE_MOVEMENT_COST * rubble_at_target) * weather_cfg["power_loss_factor"])
-    def move(self, direction, repeat=0):
+        return math.floor(self.unit_cfg.MOVE_COST + self.unit_cfg.RUBBLE_MOVEMENT_COST * rubble_at_target)
+    def move(self, direction, repeat=0, n=1):
         if isinstance(direction, int):
             direction = direction
         else:
             pass
-        return np.array([0, direction, 0, 0, repeat])
+        return np.array([0, direction, 0, 0, repeat, n])
 
-    def transfer(self, transfer_direction, transfer_resource, transfer_amount, repeat=0):
+    def transfer(self, transfer_direction, transfer_resource, transfer_amount, repeat=0, n=1):
         assert transfer_resource < 5 and transfer_resource >= 0
         assert transfer_direction < 5 and transfer_direction >= 0
-        return np.array([1, transfer_direction, transfer_resource, transfer_amount, repeat])
+        return np.array([1, transfer_direction, transfer_resource, transfer_amount, repeat, n])
     
-    def pickup(self, pickup_resource, pickup_amount, repeat=0):
+    def pickup(self, pickup_resource, pickup_amount, repeat=0, n=1):
         assert pickup_resource < 5 and pickup_resource >= 0
-        return np.array([2, 0, pickup_resource, pickup_amount, repeat])
+        return np.array([2, 0, pickup_resource, pickup_amount, repeat, n])
     
     def dig_cost(self, game_state):
-        current_weather = game_state.weather_schedule[game_state.real_env_steps]
-        weather_cfg = get_weather_config(current_weather, self.env_cfg)
-        return math.ceil(self.unit_cfg.DIG_COST * weather_cfg["power_loss_factor"])
-    def dig(self, repeat=0):
-        return np.array([3, 0, 0, 0, repeat])
+        return self.unit_cfg.DIG_COST
+    def dig(self, repeat=0, n=1):
+        return np.array([3, 0, 0, 0, repeat, n])
 
     def self_destruct_cost(self, game_state):
-        current_weather = game_state.weather_schedule[game_state.real_env_steps]
-        weather_cfg = get_weather_config(current_weather, self.env_cfg)
-        return math.ceil(self.unit_cfg.SELF_DESTRUCT_COST * weather_cfg["power_loss_factor"])
-    def self_destruct(self, repeat=0):
-        return np.array([4, 0, 0, 0, repeat])
+        return self.unit_cfg.SELF_DESTRUCT_COST
+    def self_destruct(self, repeat=0, n=1):
+        return np.array([4, 0, 0, 0, repeat, n])
 
-    def recharge(self, x, repeat=0):
-        return np.array([5, 0, 0, x, repeat])
+    def recharge(self, x, repeat=0, n=1):
+        return np.array([5, 0, 0, x, repeat, n])
 
     def __str__(self) -> str:
         out = f"[{self.team_id}] {self.unit_id} {self.unit_type} at {self.pos}"
diff --git a/kits/python/lux/weather.py b/kits/python/lux/weather.py
deleted file mode 100644
index cc5ee749..00000000
--- a/kits/python/lux/weather.py
+++ /dev/null
@@ -1,15 +0,0 @@
-import sys
-import numpy as np
-
-def get_weather_config(current_weather, env_cfg):
-    if type(current_weather) == int or isinstance(current_weather, (int, np.integer)):
-        current_weather = env_cfg.WEATHER_ID_TO_NAME[current_weather]
-    if current_weather == "MARS_QUAKE":
-        return dict(power_gain_factor=1, power_loss_factor=1)
-    elif current_weather == "COLD_SNAP":
-        return dict(power_gain_factor=1, power_loss_factor=env_cfg.WEATHER["COLD_SNAP"]["POWER_CONSUMPTION"])
-    elif current_weather == "DUST_STORM":
-        return dict(power_gain_factor=env_cfg.WEATHER["DUST_STORM"]["POWER_GAIN"], power_loss_factor=1)
-    elif current_weather == "SOLAR_FLARE":
-        return dict(power_gain_factor=env_cfg.WEATHER["SOLAR_FLARE"]["POWER_GAIN"], power_loss_factor=1)
-    return dict(power_gain_factor=1, power_loss_factor=1)
\ No newline at end of file
diff --git a/kits/sample_env_cfg.json b/kits/sample_env_cfg.json
index 42a68b9b..157ce3fd 100644
--- a/kits/sample_env_cfg.json
+++ b/kits/sample_env_cfg.json
@@ -17,10 +17,6 @@
   "MAX_RUBBLE": 100,
   "MIN_FACTORIES": 2,
   "MIN_LICHEN_TO_SPREAD": 20,
-  "NUM_WEATHER_EVENTS_RANGE": [
-    3,
-    5
-  ],
   "ORE_METAL_RATIO": 5,
   "ROBOTS": {
     "HEAVY": {
@@ -54,47 +50,11 @@
       "MOVE_COST": 1,
       "POWER_COST": 50,
       "RUBBLE_AFTER_DESTRUCTION": 1,
-      "RUBBLE_MOVEMENT_COST": 1,
+      "RUBBLE_MOVEMENT_COST": 0.1,
       "SELF_DESTRUCT_COST": 5
     }
   },
   "UNIT_ACTION_QUEUE_SIZE": 20,
-  "WEATHER": {
-    "COLD_SNAP": {
-      "POWER_CONSUMPTION": 2,
-      "TIME_RANGE": [
-        10,
-        30
-      ]
-    },
-    "DUST_STORM": {
-      "POWER_CONSUMPTION": 0.5,
-      "TIME_RANGE": [
-        10,
-        30
-      ]
-    },
-    "MARS_QUAKE": {
-      "TIME_RANGE": [
-        1,
-        5
-      ]
-    },
-    "SOLAR_FLARE": {
-      "POWER_GAIN": 2,
-      "TIME_RANGE": [
-        10,
-        30
-      ]
-    }
-  },
-  "WEATHER_ID_TO_NAME": [
-    "NONE",
-    "MARS_QUAKE",
-    "COLD_SNAP",
-    "DUST_STORM",
-    "SOLAR_FLARE"
-  ],
   "map_size": 48,
   "max_episode_length": 1000,
   "max_transfer_amount": 10000,
diff --git a/kits/sample_obs.json b/kits/sample_obs.json
index 004f04ef..88c8c6d5 100644
--- a/kits/sample_obs.json
+++ b/kits/sample_obs.json
@@ -3389,47 +3389,5 @@
     },
     "factories_per_team": 4
   },
-  "weather_schedule": [
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
-    3, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
-    3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
-  ],
   "real_env_steps": 10
 }
diff --git a/luxai_s2/luxai_runner/bot.py b/luxai_s2/luxai_runner/bot.py
index e303a3c9..82f7ffcc 100644
--- a/luxai_s2/luxai_runner/bot.py
+++ b/luxai_s2/luxai_runner/bot.py
@@ -1,17 +1,25 @@
-from argparse import Namespace
 import asyncio
 import json
-import os.path as osp
 import os
-from subprocess import Popen
+import os.path as osp
 import time
+from argparse import Namespace
+from subprocess import Popen
+
 from luxai_runner.ext_to_command import ext_to_command
 from luxai_runner.logger import Logger
 from luxai_runner.process import BotProcess
 
 
 class Bot:
-    def __init__(self, main_file_path: str, agent: str, agent_idx: str, verbose: int = 1, direct_import_python_bots = False) -> None:
+    def __init__(
+        self,
+        main_file_path: str,
+        agent: str,
+        agent_idx: str,
+        verbose: int = 1,
+        direct_import_python_bots=False,
+    ) -> None:
         """
         if direct_import_python_bots is True, will directly import the python agents and call their agent_fn functions.
         """
@@ -29,17 +37,34 @@ def __init__(self, main_file_path: str, agent: str, agent_idx: str, verbose: int
         self.agent_idx = agent_idx
         self.verbose = verbose
         self.direct_import_python_bots = direct_import_python_bots
-        self.proc = BotProcess(self.command, self.main_file_path, verbose=verbose, direct_import_python_bots=direct_import_python_bots)
+        self.proc = BotProcess(
+            self.command,
+            self.main_file_path,
+            verbose=verbose,
+            direct_import_python_bots=direct_import_python_bots,
+        )
         # timing
         self.remainingOverageTime = 60
         self.time_per_step = 3
 
-        self.log = Logger(identifier=f"{self.agent}, {self.main_file_path}",verbosity=verbose)
+        self.log = Logger(
+            identifier=f"{self.agent}, {self.main_file_path}", verbosity=verbose
+        )
 
     async def step(self, obs, step: int, reward: float = 0, info=dict()):
         stime = time.time()
         import copy
-        observations = copy.deepcopy(dict(obs=obs, step=step, remainingOverageTime=self.remainingOverageTime, player=self.agent, reward=float(reward), info=info))
+
+        observations = copy.deepcopy(
+            dict(
+                obs=obs,
+                step=step,
+                remainingOverageTime=self.remainingOverageTime,
+                player=self.agent,
+                reward=float(reward),
+                info=info,
+            )
+        )
         stderr = None
         try:
             if self.direct_import_python_bots and self.command == "python":
@@ -51,7 +76,10 @@ async def step(self, obs, step: int, reward: float = 0, info=dict()):
                 action = self.proc.agent_fn(observations, dict(env_cfg=env_cfg))
             else:
                 data = json.dumps(observations)
-                action, stderr = await asyncio.wait_for(self.proc.write(f"{data}\n"), timeout=self.remainingOverageTime + self.time_per_step)
+                action, stderr = await asyncio.wait_for(
+                    self.proc.write(f"{data}\n"),
+                    timeout=self.remainingOverageTime + self.time_per_step,
+                )
         except asyncio.TimeoutError:
             action, stderr = None, None
         time_used = time.time() - stime
@@ -68,7 +96,8 @@ async def step(self, obs, step: int, reward: float = 0, info=dict()):
             action = None
         else:
             try:
-                if isinstance(action, dict): return action
+                if isinstance(action, dict):
+                    return action
                 action = json.loads(action)
                 if not isinstance(action, dict):
                     raise ValueError("")
@@ -76,4 +105,3 @@ async def step(self, obs, step: int, reward: float = 0, info=dict()):
                 self.log.err(f"cannot parse action '{action}'")
                 action = None
         return action
-
diff --git a/luxai_s2/luxai_runner/cli.py b/luxai_s2/luxai_runner/cli.py
index f884a723..70cdafb9 100644
--- a/luxai_s2/luxai_runner/cli.py
+++ b/luxai_s2/luxai_runner/cli.py
@@ -1,72 +1,126 @@
 import asyncio
+import json
+import sys
 from typing import Dict, List
-from luxai_runner.bot import Bot
-from luxai_s2 import LuxAI_S2
+
 import numpy as np
-import json
+from luxai_runner.bot import Bot
 from luxai_runner.episode import Episode, EpisodeConfig, ReplayConfig
-from luxai_runner.tournament import Tournament, TournamentConfig
 from luxai_runner.logger import Logger
+from luxai_runner.tournament import Tournament, TournamentConfig
 from omegaconf import OmegaConf
-import sys
+
+from luxai_s2 import LuxAI_S2
+
 
 def main():
     np.random.seed(0)
     import argparse
 
     parser = argparse.ArgumentParser(description="Run the LuxAI 2022 game.")
-    parser.add_argument("players", nargs="+", help="Paths to player modules. If --tournament is passed as well, you can also pass a folder and we will look through all sub-folders for valid agents with main.py files (only works for python agents at the moment).")
-    parser.add_argument("-l", "--len", help="Max episode length", type=int, default=1000)
+    parser.add_argument(
+        "players",
+        nargs="+",
+        help="Paths to player modules. If --tournament is passed as well, you can also pass a folder and we will look through all sub-folders for valid agents with main.py files (only works for python agents at the moment).",
+    )
+    parser.add_argument(
+        "-l", "--len", help="Max episode length", type=int, default=1000
+    )
 
     # replay configs
-    parser.add_argument("-o", "--output", help="Where to output replays. Default is none and no replay is generated")
-    parser.add_argument("--replay.save_format", help="Save format \"json\" works with the visualizer while pickle is a compact, python usable version", default="json")
-    parser.add_argument("--replay.compressed_obs", help="Whether to save compressed observations or not. Compressed observations do not contain the full observation at each step. In particular, the map information is stored as the first observation, subsequent observations only store the changes that happened.", default=True)
+    parser.add_argument(
+        "-o",
+        "--output",
+        help="Where to output replays. Default is none and no replay is generated",
+    )
+    parser.add_argument(
+        "--replay.save_format",
+        help='Save format "json" works with the visualizer while pickle is a compact, python usable version',
+        default="json",
+    )
+    parser.add_argument(
+        "--replay.compressed_obs",
+        help="Whether to save compressed observations or not. Compressed observations do not contain the full observation at each step. In particular, the map information is stored as the first observation, subsequent observations only store the changes that happened.",
+        default=True,
+    )
 
     # episode configs
     parser.add_argument(
-        "-v", "--verbose", help="Verbose Level (0 = silent, 1 = errors, 2 = warnings, 3 = info)", type=int, default=1
+        "-v",
+        "--verbose",
+        help="Verbose Level (0 = silent, 1 = errors, 2 = warnings, 3 = info)",
+        type=int,
+        default=1,
+    )
+    parser.add_argument(
+        "-s",
+        "--seed",
+        help="Fix a seed for episode(s). All episodes will initialize the same, including tournament ones",
+        type=int,
     )
-    parser.add_argument("-s", "--seed", help="Fix a seed for episode(s). All episodes will initialize the same, including tournament ones", type=int)
 
     # env configs
 
-    parser.add_argument("--render", help="Render with a window", action="store_true", default=False)
-
-    parser.add_argument("--tournament", help="Turn tournament mode on", action="store_true", default=False)
-    parser.add_argument("--tournament_cfg.concurrent", help="Max concurrent number of episodes to run. Recommended to set no higher than the number of CPUs / 2", type=int, default=1)
-    parser.add_argument("--tournament_cfg.ranking_system", help="The ranking system to use. Default is 'elo'. Can be 'elo', 'wins'.", type=str, default="elo")
-    parser.add_argument("--skip_validate_action_space", help="Set this for a small performance increase. Note that turning this on means the engine assumes your submitted actions are valid. If your actions are not well formatted there could be errors", action="store_true", default=False)
+    parser.add_argument(
+        "--render", help="Render with a window", action="store_true", default=False
+    )
 
+    parser.add_argument(
+        "--tournament",
+        help="Turn tournament mode on",
+        action="store_true",
+        default=False,
+    )
+    parser.add_argument(
+        "--tournament_cfg.concurrent",
+        help="Max concurrent number of episodes to run. Recommended to set no higher than the number of CPUs / 2",
+        type=int,
+        default=1,
+    )
+    parser.add_argument(
+        "--tournament_cfg.ranking_system",
+        help="The ranking system to use. Default is 'elo'. Can be 'elo', 'wins'.",
+        type=str,
+        default="elo",
+    )
+    parser.add_argument(
+        "--skip_validate_action_space",
+        help="Set this for a small performance increase. Note that turning this on means the engine assumes your submitted actions are valid. If your actions are not well formatted there could be errors",
+        action="store_true",
+        default=False,
+    )
 
     args = parser.parse_args()
 
-    if sys.platform == 'win32':
+    if sys.platform == "win32":
         asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
-    
+
     cfg = EpisodeConfig(
-            players=args.players,
-            env_cls=LuxAI_S2,
-            seed=args.seed,
-            env_cfg=dict(
-                verbose=args.verbose,
-                validate_action_space=not args.skip_validate_action_space,
-                max_episode_length=args.len,
-            ),
-            verbosity=args.verbose,
-            save_replay_path=args.output,
-            replay_options=ReplayConfig(
-                save_format=getattr(args, "replay.save_format"),
-                compressed_obs=getattr(args, "replay.compressed_obs")
-            ),
-            render=args.render
-        )
+        players=args.players,
+        env_cls=LuxAI_S2,
+        seed=args.seed,
+        env_cfg=dict(
+            verbose=args.verbose,
+            validate_action_space=not args.skip_validate_action_space,
+            max_episode_length=args.len,
+        ),
+        verbosity=args.verbose,
+        save_replay_path=args.output,
+        replay_options=ReplayConfig(
+            save_format=getattr(args, "replay.save_format"),
+            compressed_obs=getattr(args, "replay.compressed_obs"),
+        ),
+        render=args.render,
+    )
 
     if args.tournament:
         import os
+
         if os.path.isdir(args.players[0]):
-            assert len(args.players) == 1, "Found more than one positional argument despite being given a directory of players"
-            subfolders = [ f.path for f in os.scandir(args.players[0]) if f.is_dir() ]
+            assert (
+                len(args.players) == 1
+            ), "Found more than one positional argument despite being given a directory of players"
+            subfolders = [f.path for f in os.scandir(args.players[0]) if f.is_dir()]
             agents = []
             for sub_dir in subfolders:
                 agent_file = os.path.join(sub_dir, "main.py")
@@ -78,24 +132,27 @@ def main():
         tournament_config = TournamentConfig()
         tournament_config.agents = args.players
         # TODO - in future replace this with OmegaConf or something that can parse these nicely
-        tournament_config.max_concurrent_episodes = getattr(args, "tournament_cfg.concurrent")
-        tournament_config.ranking_system = getattr(args, "tournament_cfg.ranking_system")
+        tournament_config.max_concurrent_episodes = getattr(
+            args, "tournament_cfg.concurrent"
+        )
+        tournament_config.ranking_system = getattr(
+            args, "tournament_cfg.ranking_system"
+        )
         tourney = Tournament(
-            cfg=tournament_config,
-            episode_cfg=cfg # the base/default episode config
+            cfg=tournament_config, episode_cfg=cfg  # the base/default episode config
         )
         # import ipdb;ipdb.set_trace()
         asyncio.run(tourney.run())
         # exit()
     else:
         import time
+
         stime = time.time()
-        eps = Episode(
-            cfg=cfg
-        )
+        eps = Episode(cfg=cfg)
         asyncio.run(eps.run())
         etime = time.time()
         print(etime - stime)
 
+
 if __name__ == "__main__":
-    main()
\ No newline at end of file
+    main()
diff --git a/luxai_s2/luxai_runner/episode.py b/luxai_s2/luxai_runner/episode.py
index 4de6dce7..6e4569ec 100644
--- a/luxai_s2/luxai_runner/episode.py
+++ b/luxai_s2/luxai_runner/episode.py
@@ -1,21 +1,24 @@
-
 import asyncio
-from dataclasses import dataclass
+import dataclasses
 import json
+import os.path as osp
 import time
+from dataclasses import dataclass
 from typing import Any, Callable, Dict, List, Optional
+
 import gym
-from luxai_runner.logger import Logger
-from luxai_runner.bot import Bot
 import numpy as np
-import dataclasses
+from luxai_runner.bot import Bot
+from luxai_runner.logger import Logger
 from luxai_runner.utils import to_json
-import os.path as osp
 
-@dataclass 
+
+@dataclass
 class ReplayConfig:
     save_format: str = "json"
     compressed_obs: bool = False
+
+
 @dataclass
 class EpisodeConfig:
     players: List[str]
@@ -27,6 +30,7 @@ class EpisodeConfig:
     save_replay_path: Optional[str] = None
     replay_options: ReplayConfig = ReplayConfig()
 
+
 class Episode:
     def __init__(self, cfg: EpisodeConfig) -> None:
         self.cfg = cfg
@@ -44,6 +48,7 @@ def save_replay(self, replay):
             del replay["rewards"]
             ext = ".json"
             from pathlib import Path
+
             dir_name = osp.dirname(self.cfg.save_replay_path)
             if dir_name != "":
                 Path(dir_name).mkdir(parents=True, exist_ok=True)
@@ -52,10 +57,12 @@ def save_replay(self, replay):
             with open(f"{self.cfg.save_replay_path}{ext}", "w") as f:
                 json.dump(replay, f)
         else:
-            raise ValueError(f"{self.cfg.replay_options.save_format} is not a valid save format")
+            raise ValueError(
+                f"{self.cfg.replay_options.save_format} is not a valid save format"
+            )
 
     async def run(self):
-        if len(self.players) != 2: 
+        if len(self.players) != 2:
             raise ValueError("Must provide two paths.")
         # Start agents
         players: Dict[str, Bot] = dict()
@@ -68,19 +75,18 @@ async def run(self):
             start_tasks += [player.proc.start()]
         await asyncio.wait(start_tasks, return_when=asyncio.ALL_COMPLETED)
 
-        
         obs = self.env.reset(seed=self.seed)
         env_cfg = self.env.state.env_cfg
         state_obs = self.env.state.get_compressed_obs()
         obs = to_json(state_obs)
 
-        if self.cfg.render: 
+        if self.cfg.render:
             self.env.render()
             time.sleep(0.2)
         game_done = False
         rewards, dones, infos = dict(), dict(), dict()
         for agent in players:
-            rewards[agent] = 0 
+            rewards[agent] = 0
             dones[agent] = 0
             infos[agent] = dict(
                 # turn 0 provide configurations
@@ -99,11 +105,13 @@ async def run(self):
             i += 1
             # print("===", self.env.env_steps)
             actions = dict()
-            
+
             agent_ids = []
             action_coros = []
             for player in players.values():
-                action = player.step(obs, self.env.env_steps, rewards[agent], infos[agent])
+                action = player.step(
+                    obs, self.env.env_steps, rewards[agent], infos[agent]
+                )
                 action_coros += [action]
                 agent_ids += [player.agent]
             resolved_actions = await asyncio.gather(*action_coros)
@@ -116,7 +124,9 @@ async def run(self):
                 except:
                     if self.cfg.verbosity > 0:
                         if action is None:
-                            print(f"{agent_id} sent a invalid action {action}. Agent likely errored out somewhere, check above for stderr logs")
+                            print(
+                                f"{agent_id} sent a invalid action {action}. Agent likely errored out somewhere, check above for stderr logs"
+                            )
                         else:
                             print(f"{agent_id} sent a invalid action {action}")
                     actions[agent_id] = None
@@ -133,13 +143,14 @@ async def run(self):
                 replay["rewards"].append(rewards)
                 replay["dones"].append(dones)
 
-            if self.cfg.render: 
+            if self.cfg.render:
                 self.env.render()
                 time.sleep(0.1)
             players_left = len(dones)
             for k in dones:
-                if dones[k]: players_left -= 1
-            if players_left < 2: # specific to lux ai 2022
+                if dones[k]:
+                    players_left -= 1
+            if players_left < 2:  # specific to lux ai 2022
                 game_done = True
         self.log.info(f"Final Scores: {rewards}")
         if save_replay:
diff --git a/luxai_s2/luxai_runner/logger.py b/luxai_s2/luxai_runner/logger.py
index e967f2c1..8f793a56 100644
--- a/luxai_s2/luxai_runner/logger.py
+++ b/luxai_s2/luxai_runner/logger.py
@@ -1,11 +1,14 @@
 from turtle import color
 
 from luxai_s2.globals import TERM_COLORS
+
 try:
     from termcolor import colored
 except:
     pass
 import os
+
+
 class Logger:
     """
     A basic logger
@@ -16,9 +19,11 @@ class Logger:
         2 : Log warnings
         3 : Log everything
     """
+
     def __init__(self, identifier: str, verbosity: int = 1) -> None:
         self.verbosity = verbosity
         self.identifier = identifier
+
     def _print(self, msg: str, color: str, end: str = "\n"):
         if self.identifier != "":
             msg = f"{self.identifier}: {msg}"
@@ -26,9 +31,15 @@ def _print(self, msg: str, color: str, end: str = "\n"):
             print(colored(msg, color), end=end)
         else:
             print(msg, end=end)
+
     def err(self, msg: str, end: str = "\n"):
-        if self.verbosity >= 1: self._print(msg, "red", end=end)
+        if self.verbosity >= 1:
+            self._print(msg, "red", end=end)
+
     def warn(self, msg: str, end: str = "\n"):
-        if self.verbosity >= 2: self._print(msg, "yellow", end=end)
+        if self.verbosity >= 2:
+            self._print(msg, "yellow", end=end)
+
     def info(self, msg: str, end: str = "\n"):
-        if self.verbosity >= 3: self._print(msg, "cyan", end=end)
\ No newline at end of file
+        if self.verbosity >= 3:
+            self._print(msg, "cyan", end=end)
diff --git a/luxai_s2/luxai_runner/process.py b/luxai_s2/luxai_runner/process.py
index 380e30df..d759ef20 100644
--- a/luxai_s2/luxai_runner/process.py
+++ b/luxai_s2/luxai_runner/process.py
@@ -1,18 +1,25 @@
-from datetime import datetime
-from subprocess import Popen, PIPE, STDOUT
-from threading import Thread
-from queue import Queue, Empty
 import asyncio
-
 import atexit
 import os
 import sys
+from datetime import datetime
+from queue import Empty, Queue
+from subprocess import PIPE, STDOUT, Popen
+from threading import Thread
 from typing import IO
+
 from luxai_runner.logger import Logger
 
 
 class BotProcess:
-    def __init__(self, command: str, file_path: str, verbose: int = 2, live_log: str = True, direct_import_python_bots = False) -> None:
+    def __init__(
+        self,
+        command: str,
+        file_path: str,
+        verbose: int = 2,
+        live_log: str = True,
+        direct_import_python_bots=False,
+    ) -> None:
         self.command = command
         if self.command == "./":
             self.is_binary = True
@@ -26,6 +33,7 @@ def __init__(self, command: str, file_path: str, verbose: int = 2, live_log: str
         if self.direct_import_python_bots and self.command == "python":
             import importlib.util
             import sys
+
             sys.path.append(os.path.dirname(file_path))
             spec = importlib.util.spec_from_file_location("bot", file_path)
             foo = importlib.util.module_from_spec(spec)
@@ -35,7 +43,8 @@ def __init__(self, command: str, file_path: str, verbose: int = 2, live_log: str
             self.agent_fn = foo.agent_fn
 
     async def start(self):
-        if self.direct_import_python_bots and self.command == "python": return
+        if self.direct_import_python_bots and self.command == "python":
+            return
         cwd = os.path.dirname(self.file_path)
         if cwd == "":
             cwd = "."
@@ -50,7 +59,7 @@ async def start(self):
                 stdout=asyncio.subprocess.PIPE,
                 stderr=asyncio.subprocess.PIPE,
                 cwd=cwd,
-                limit=1024 * 128
+                limit=1024 * 128,
             )
         else:
             if self.command == "java" and base_file_path.endswith(".java"):
@@ -62,7 +71,7 @@ async def start(self):
                 stdout=asyncio.subprocess.PIPE,
                 stderr=asyncio.subprocess.PIPE,
                 cwd=cwd,
-                limit=1024 * 128
+                limit=1024 * 128,
             )
         self.log.info(f"Started {self.command} {self.file_path}")
         # following 4 lines from https://stackoverflow.com/questions/375427/a-non-blocking-read-on-a-subprocess-pipe-in-python
@@ -77,20 +86,27 @@ def enqueue_output(out: IO, queue: Queue):
         # self._t.daemon = True  # thread dies with the program
         # self._t.start()
         self.stderr_queue = []
+
         async def log_stream(stream):
             while not stream.at_eof():
                 data = await stream.readline()
                 line = data.decode()
 
-                if stream.at_eof() and len(line) == 0: break
-                elif self.live_log: self.log.err(line, end="")
-                else: self.stderr_queue.append(line)
+                if stream.at_eof() and len(line) == 0:
+                    break
+                elif self.live_log:
+                    self.log.err(line, end="")
+                else:
+                    self.stderr_queue.append(line)
+
         asyncio.create_task(log_stream(self._agent_process.stderr))
         # await asyncio.gather(watch(self._agent_process.stderr, 'E:'))
 
     async def write(self, msg: str):
         self._agent_process.stdin.write(msg.encode())
-        stdout, stderr = await asyncio.gather(self._agent_process.stdout.readline(), self.stderr())
+        stdout, stderr = await asyncio.gather(
+            self._agent_process.stdout.readline(), self.stderr()
+        )
         return stdout.decode(), stderr
 
     async def receive(self) -> str:
@@ -98,7 +114,7 @@ async def receive(self) -> str:
         return res
 
     async def stderr(self):
-        r =  "".join(self.stderr_queue)
+        r = "".join(self.stderr_queue)
         self.stderr_queue.clear()
         return r
         # while True:
diff --git a/luxai_s2/luxai_runner/tournament/__init__.py b/luxai_s2/luxai_runner/tournament/__init__.py
index 4647eb4c..b1c232c1 100644
--- a/luxai_s2/luxai_runner/tournament/__init__.py
+++ b/luxai_s2/luxai_runner/tournament/__init__.py
@@ -1,2 +1,2 @@
 from .config import TournamentConfig
-from .tournament import Tournament
\ No newline at end of file
+from .tournament import Tournament
diff --git a/luxai_s2/luxai_runner/tournament/config.py b/luxai_s2/luxai_runner/tournament/config.py
index 2b94e393..59425235 100644
--- a/luxai_s2/luxai_runner/tournament/config.py
+++ b/luxai_s2/luxai_runner/tournament/config.py
@@ -5,8 +5,8 @@
 @dataclass
 class TournamentConfig:
     name: str = "luxai_s2_tourney"
-    tournament_type: str = "leadeboard" # roundrobin, leaderboard
-    ranking_system: str = "elo" # elo, win/loss, TODO: Bradley-Terry
+    tournament_type: str = "leadeboard"  # roundrobin, leaderboard
+    ranking_system: str = "elo"  # elo, win/loss, TODO: Bradley-Terry
     ranking_system_cfg: Dict = field(default_factory=dict)
     matchmaking_system: str = "random"
 
@@ -14,4 +14,4 @@ class TournamentConfig:
 
     agents_per_episode: List[int] = field(default_factory=lambda: [2])
 
-    max_concurrent_episodes: int = 4
\ No newline at end of file
+    max_concurrent_episodes: int = 4
diff --git a/luxai_s2/luxai_runner/tournament/matchmaking.py b/luxai_s2/luxai_runner/tournament/matchmaking.py
index ffb7dc6c..7951be4f 100644
--- a/luxai_s2/luxai_runner/tournament/matchmaking.py
+++ b/luxai_s2/luxai_runner/tournament/matchmaking.py
@@ -1,22 +1,29 @@
 from abc import ABC
 from typing import List
+
 import numpy as np
 
+
 class MatchMakingSystem(ABC):
-    def __init__(self, agents: List[str], agents_per_episode: List[int], seed=0) -> None:
+    def __init__(
+        self, agents: List[str], agents_per_episode: List[int], seed=0
+    ) -> None:
         super().__init__()
         self.agents = agents
         self.agents_per_episode = agents_per_episode
         self.np_random = np.random.RandomState(seed=seed)
+
     def next_match(self) -> List[str]:
         pass
-        
 
 
 class Random(MatchMakingSystem):
-    def __init__(self, agents: List[str], agents_per_episode: List[int], seed=0) -> None:
+    def __init__(
+        self, agents: List[str], agents_per_episode: List[int], seed=0
+    ) -> None:
         super().__init__(agents, agents_per_episode, seed)
+
     def next_match(self) -> List[str]:
         num_agents = self.np_random.choice(self.agents_per_episode)
         agents = self.np_random.choice(self.agents, size=num_agents, replace=False)
-        return agents
\ No newline at end of file
+        return agents
diff --git a/luxai_s2/luxai_runner/tournament/rankingsystem.py b/luxai_s2/luxai_runner/tournament/rankingsystem.py
index cfd76d43..30566eb5 100644
--- a/luxai_s2/luxai_runner/tournament/rankingsystem.py
+++ b/luxai_s2/luxai_runner/tournament/rankingsystem.py
@@ -1,46 +1,71 @@
 from abc import ABC, abstractclassmethod
 from dataclasses import dataclass
+
 import numpy as np
 
+
 class Rank:
     rating: int
 
+
 class RankingSystem(ABC):
     def __init__(self) -> None:
         pass
 
     def init_rank_state(self) -> Rank:
         pass
-    def update(self, rank_1: Rank, rank_2: Rank, rank_1_score: float, rank_2_score: float):
+
+    def update(
+        self, rank_1: Rank, rank_2: Rank, rank_1_score: float, rank_2_score: float
+    ):
         pass
+
     def _rank_headers(self) -> str:
         pass
+
     def _rank_info(self, rank: Rank) -> str:
         pass
 
+
 @dataclass
 class ELORank(Rank):
     rating: int
     episodes: int
+
+
 class ELO(RankingSystem):
-    def __init__(self, K=32, init_rating=1000,col_length=10) -> None:
+    def __init__(self, K=32, init_rating=1000, col_length=10) -> None:
         super().__init__()
         self.col_length = col_length
         self.K = K
         self.init_rating = 1000
+
     def init_rank_state(self) -> Rank:
         return ELORank(rating=self.init_rating, episodes=0)
+
     def update(self, rank_1: ELORank, rank_2: ELORank, rank_1_score, rank_2_score):
         # Only implements win/loss/ties
         if rank_1_score > rank_2_score:
-            rank_1.rating = rank_1.rating + self.K * (1 - self._expected_score(rank_1, rank_2))
-            rank_2.rating = rank_2.rating + self.K * (0 - self._expected_score(rank_2, rank_1))
+            rank_1.rating = rank_1.rating + self.K * (
+                1 - self._expected_score(rank_1, rank_2)
+            )
+            rank_2.rating = rank_2.rating + self.K * (
+                0 - self._expected_score(rank_2, rank_1)
+            )
         elif rank_1_score < rank_2_score:
-            rank_1.rating = rank_1.rating + self.K * (0 - self._expected_score(rank_1, rank_2))
-            rank_2.rating = rank_2.rating + self.K * (1 - self._expected_score(rank_2, rank_1))
+            rank_1.rating = rank_1.rating + self.K * (
+                0 - self._expected_score(rank_1, rank_2)
+            )
+            rank_2.rating = rank_2.rating + self.K * (
+                1 - self._expected_score(rank_2, rank_1)
+            )
         else:
-            rank_1.rating = rank_1.rating + self.K * (0.5 - self._expected_score(rank_1, rank_2))
-            rank_2.rating = rank_2.rating + self.K * (0.5 - self._expected_score(rank_2, rank_1))
+            rank_1.rating = rank_1.rating + self.K * (
+                0.5 - self._expected_score(rank_1, rank_2)
+            )
+            rank_2.rating = rank_2.rating + self.K * (
+                0.5 - self._expected_score(rank_2, rank_1)
+            )
         rank_1.episodes += 1
         rank_2.episodes += 1
 
@@ -52,9 +77,11 @@ def _expected_score(self, rank_1: ELORank, rank_2: ELORank):
 
     def _rank_headers(self) -> str:
         return f"{'Rating':{self.col_length}.{self.col_length}}"
+
     def _rank_info(self, rank: ELORank) -> str:
         return f"{str(rank.rating):{self.col_length}.{self.col_length}}"
 
+
 @dataclass
 class WinLossRank(Rank):
     rating: int
@@ -63,16 +90,23 @@ class WinLossRank(Rank):
     losses: int
     episodes: int
 
+
 class WinLoss(RankingSystem):
-    def __init__(self, win_points=3,tie_points=1,loss_points=0,col_length=10) -> None:
+    def __init__(
+        self, win_points=3, tie_points=1, loss_points=0, col_length=10
+    ) -> None:
         super().__init__()
-        self.col_length=col_length
-        self.win_points=win_points
-        self.tie_points=tie_points
-        self.loss_points=loss_points
+        self.col_length = col_length
+        self.win_points = win_points
+        self.tie_points = tie_points
+        self.loss_points = loss_points
+
     def init_rank_state(self) -> Rank:
-        return WinLossRank(rating=0, episodes=0, wins=0,ties=0,losses=0)
-    def update(self, rank_1: WinLossRank, rank_2: WinLossRank, rank_1_score, rank_2_score):
+        return WinLossRank(rating=0, episodes=0, wins=0, ties=0, losses=0)
+
+    def update(
+        self, rank_1: WinLossRank, rank_2: WinLossRank, rank_1_score, rank_2_score
+    ):
         # Only implements win/loss/ties
         if rank_1_score == rank_2_score:
             winner = None
@@ -95,7 +129,9 @@ def update(self, rank_1: WinLossRank, rank_2: WinLossRank, rank_1_score, rank_2_
 
         rank_1.episodes += 1
         rank_2.episodes += 1
+
     def _rank_headers(self) -> str:
         return f"{'Score':{self.col_length}} | {'Wins':{self.col_length}} | {'Ties':{self.col_length}} | {'Losses':{self.col_length}}"
+
     def _rank_info(self, rank: WinLossRank) -> str:
-        return f"{str(rank.rating):{self.col_length}.{self.col_length}} | {str(rank.wins):{self.col_length}.{self.col_length}} | {str(rank.ties):{self.col_length}.{self.col_length}} | {str(rank.losses):{self.col_length}.{self.col_length}}"
\ No newline at end of file
+        return f"{str(rank.rating):{self.col_length}.{self.col_length}} | {str(rank.wins):{self.col_length}.{self.col_length}} | {str(rank.ties):{self.col_length}.{self.col_length}} | {str(rank.losses):{self.col_length}.{self.col_length}}"
diff --git a/luxai_s2/luxai_runner/tournament/tournament.py b/luxai_s2/luxai_runner/tournament/tournament.py
index 79f59b2c..87ead511 100644
--- a/luxai_s2/luxai_runner/tournament/tournament.py
+++ b/luxai_s2/luxai_runner/tournament/tournament.py
@@ -1,18 +1,22 @@
 import asyncio
+import copy
 import os
+import os.path as osp
 from typing import Dict
+
 from luxai_runner.episode import Episode, EpisodeConfig
+from luxai_runner.tournament import matchmaking
 from luxai_runner.tournament.config import TournamentConfig
 from luxai_runner.tournament.rankingsystem import ELO, Rank, WinLoss, WinLossRank
-from luxai_runner.tournament import matchmaking
-import copy
-import os.path as osp
+
+
 class Player:
     def __init__(self, player_id, file) -> None:
         self.id = player_id
         self.file = file
         self.rank: Rank = None
 
+
 class Tournament:
     def __init__(self, cfg: TournamentConfig, episode_cfg: EpisodeConfig = None):
         self.global_id = 0
@@ -22,11 +26,11 @@ def __init__(self, cfg: TournamentConfig, episode_cfg: EpisodeConfig = None):
         min_agents = min(self.cfg.agents_per_episode)
         assert len(self.cfg.agents) >= min_agents
         assert episode_cfg is not None
-        
+
         # init ranking system
         if cfg.ranking_system.lower() == "elo":
             # TODO - allow configuration via CLI
-            self.ranking_sys = ELO(K=30,init_rating=1000)
+            self.ranking_sys = ELO(K=30, init_rating=1000)
         elif cfg.ranking_system.lower() == "wins":
             self.ranking_sys = WinLoss(win_points=3, tie_points=1, loss_points=0)
         else:
@@ -35,22 +39,27 @@ def __init__(self, cfg: TournamentConfig, episode_cfg: EpisodeConfig = None):
         self.players: Dict[str, Player] = dict()
         for agent in self.cfg.agents:
             self.add_player(agent)
-        
-        self.match_making_sys = matchmaking.Random([x for x in self.players], self.cfg.agents_per_episode)
 
-    def add_player(self, file, name = None):
-        if name is None: name = file
+        self.match_making_sys = matchmaking.Random(
+            [x for x in self.players], self.cfg.agents_per_episode
+        )
+
+    def add_player(self, file, name=None):
+        if name is None:
+            name = file
         name = f"{name}_{self.global_id}"
         assert os.path.isfile(file), f"Agent file {file} does not exist"
 
         self.global_id += 1
         self.players[name] = Player(name, file)
         self.players[name].rank = self.ranking_sys.init_rank_state()
-    
+
     async def run(self):
         episodes = set()
+
         async def _run_episode_cb(a):
-            if a in episodes: episodes.discard(a)
+            if a in episodes:
+                episodes.discard(a)
             next_players = self.match_making_sys.next_match()
             eps_cfg = copy.deepcopy(self.eps_cfg)
             if self.eps_cfg.save_replay_path is not None:
@@ -59,7 +68,7 @@ async def _run_episode_cb(a):
                 eps_cfg.save_replay_path = f"{save_replay_path_split[0]}_{self.episode_id}{save_replay_path_split[1]}"
             eps_cfg.players = [self.players[p].file for p in next_players]
             self.episode_id += 1
-            
+
             players = dict()
             for i in range(len(next_players)):
                 players[f"player_{i}"] = self.players[next_players[i]]
@@ -68,8 +77,10 @@ async def _run_episode_cb(a):
             episodes.add(task)
             await task
             await _run_episode_cb(task)
+
         async def print_results():
             import time
+
             line_length = 50
             line_length += len(self.ranking_sys._rank_headers())
             while True:
@@ -81,32 +92,40 @@ async def print_results():
 
                 lines.append(f"==== {self.cfg.name} ====")
                 # lines.append("")
-                lines.append(f"{'Player':36.36}| {self.ranking_sys._rank_headers()}| {'Episodes':14.14}")
-                lines.append("-"*line_length)
+                lines.append(
+                    f"{'Player':36.36}| {self.ranking_sys._rank_headers()}| {'Episodes':14.14}"
+                )
+                lines.append("-" * line_length)
                 players_sorted = [self.players[p] for p in self.players]
-                players_sorted = sorted(players_sorted, key=lambda p : p.rank.rating, reverse=True)
+                players_sorted = sorted(
+                    players_sorted, key=lambda p: p.rank.rating, reverse=True
+                )
                 for p in players_sorted:
                     rank = p.rank
-                    lines.append(f"{p.id:36.36}| {self.ranking_sys._rank_info(rank)}| {str(rank.episodes):14.14}")
-                lines.append("-"*line_length)
+                    lines.append(
+                        f"{p.id:36.36}| {self.ranking_sys._rank_info(rank)}| {str(rank.episodes):14.14}"
+                    )
+                lines.append("-" * line_length)
                 lines.append(f"{len(episodes)} episodes are running")
 
-
                 for _ in range(len(lines)):
                     sys.stdout.write("\x1b[1A\x1b[2K")
                 for i in range(len(lines)):
                     sys.stdout.write(lines[i] + "\n")
                 await asyncio.sleep(1)
+
         coros = []
         for i in range(self.cfg.max_concurrent_episodes):
             coros += [_run_episode_cb(None)]
         await asyncio.gather(*coros, print_results())
 
     async def _run_episode(self, players: Dict[str, Player], eps_cfg: EpisodeConfig):
-        eps = Episode(
-            cfg=eps_cfg
-        )
+        eps = Episode(cfg=eps_cfg)
         rewards = await eps.run()
         # TODO: there is probably some race condition here, but in long run it's not important
-        self.ranking_sys.update(players["player_0"].rank, players["player_1"].rank, rewards["player_0"], rewards["player_1"])
-        
\ No newline at end of file
+        self.ranking_sys.update(
+            players["player_0"].rank,
+            players["player_1"].rank,
+            rewards["player_0"],
+            rewards["player_1"],
+        )
diff --git a/luxai_s2/luxai_s2/__init__.py b/luxai_s2/luxai_s2/__init__.py
index 71552a47..6bfc8b7b 100644
--- a/luxai_s2/luxai_s2/__init__.py
+++ b/luxai_s2/luxai_s2/__init__.py
@@ -1,1 +1,2 @@
-from .env import LuxAI_S2
\ No newline at end of file
+from .env import LuxAI_S2
+from .version import __version__
\ No newline at end of file
diff --git a/luxai_s2/luxai_s2/actions.py b/luxai_s2/luxai_s2/actions.py
index 77fd86aa..60b1f134 100644
--- a/luxai_s2/luxai_s2/actions.py
+++ b/luxai_s2/luxai_s2/actions.py
@@ -1,33 +1,31 @@
+import math
 from collections import defaultdict
 from dataclasses import dataclass
 from enum import Enum
-import math
-from typing import Callable
+from typing import TYPE_CHECKING, Callable
 
 import numpy as np
-from typing import TYPE_CHECKING
+
 from luxai_s2.config import EnvConfig
+
 if TYPE_CHECKING:
     from luxai_s2.factory import Factory
     from luxai_s2.state import State
-from luxai_s2.map.position import Position
-
 
 import luxai_s2.unit as luxai_unit
+from luxai_s2.map.position import Position
 
 
-# (0 = move, 1 = transfer X amount of R, 2 = pickup X amount of R, 3 = dig, 4 = self destruct, 5 = recharge X, 6 = repeat)
+# (0 = move, 1 = transfer X amount of R, 2 = pickup X amount of R, 3 = dig, 4 = self destruct, 5 = recharge X)
 class Action:
     def __init__(self, act_type: str) -> None:
         self.act_type = act_type
-        self.repeat = 0
+        self.n = 1  # number of times to execute the action
         self.power_cost = 0
+        self.repeat = False  # whether to put action to the back of the action queue
 
     def state_dict(self):
         raise NotImplementedError("")
-    @property
-    def repeating(self):
-        return self.repeat == -1 or self.repeat > 0
 
 
 class FactoryBuildAction(Action):
@@ -47,79 +45,106 @@ def __init__(self) -> None:
         super().__init__("factory_water")
         self.water_cost = None
         self.power_cost = 0
+
     def state_dict(self):
         return 2
 
 
 class MoveAction(Action):
-    def __init__(self, move_dir: int, dist: int = 1, repeat=0) -> None:
+    def __init__(
+        self, move_dir: int, dist: int = 1, repeat: bool = False, n: int = 1
+    ) -> None:
         super().__init__("move")
         # a[1] = direction (0 = center, 1 = up, 2 = right, 3 = down, 4 = left)
         self.move_dir = move_dir
         self.dist = dist
         self.repeat = repeat
+        self.n = n
         self.power_cost = 0
 
     def state_dict(self):
-        return np.array([0, self.move_dir, 0, 0, self.repeat])
+        return np.array([0, self.move_dir, 0, 0, self.repeat, self.n])
 
 
 class TransferAction(Action):
-    def __init__(self, transfer_dir: int, resource: int, transfer_amount: int, repeat=0) -> None:
+    def __init__(
+        self,
+        transfer_dir: int,
+        resource: int,
+        transfer_amount: int,
+        repeat: bool = False,
+        n: int = 1,
+    ) -> None:
         super().__init__("transfer")
         # a[2] = R = resource type (0 = ice, 1 = ore, 2 = water, 3 = metal, 4 power)
         self.transfer_dir = transfer_dir
         self.resource = resource
         self.transfer_amount = transfer_amount
         self.repeat = repeat
+        self.n = n
         self.power_cost = 0
 
     def state_dict(self):
-        return np.array([1, self.transfer_dir, self.resource, self.transfer_amount, self.repeat])
+        return np.array(
+            [
+                1,
+                self.transfer_dir,
+                self.resource,
+                self.transfer_amount,
+                self.repeat,
+                self.n,
+            ]
+        )
 
 
 class PickupAction(Action):
-    def __init__(self, resource: int, pickup_amount: int, repeat=0) -> None:
+    def __init__(
+        self, resource: int, pickup_amount: int, repeat: bool = False, n: int = 1
+    ) -> None:
         super().__init__("pickup")
         # a[2] = R = resource type (0 = ice, 1 = ore, 2 = water, 3 = metal, 4 power)
         self.resource = resource
         self.pickup_amount = pickup_amount
         self.repeat = repeat
+        self.n = n
         self.power_cost = 0
 
     def state_dict(self):
-        return np.array([2, 0, self.resource, self.pickup_amount, self.repeat])
+        return np.array([2, 0, self.resource, self.pickup_amount, self.repeat, self.n])
 
 
 class DigAction(Action):
-    def __init__(self, repeat=0) -> None:
+    def __init__(self, repeat: bool = False, n: int = 1) -> None:
         super().__init__("dig")
         self.repeat = repeat
+        self.n = n
         self.power_cost = 0
 
     def state_dict(self):
-        return np.array([3, 0, 0, 0, self.repeat])
+        return np.array([3, 0, 0, 0, self.repeat, self.n])
 
 
 class SelfDestructAction(Action):
-    def __init__(self, repeat=0) -> None:
+    def __init__(self, repeat: bool = False, n: int = 1) -> None:
         super().__init__("self_destruct")
         self.repeat = repeat
+        self.n = n
         self.power_cost = 0
 
     def state_dict(self):
-        return np.array([4, 0, 0, 0, self.repeat])
+        return np.array([4, 0, 0, 0, self.repeat, self.n])
 
 
 class RechargeAction(Action):
-    def __init__(self, power: int, repeat=0) -> None:
+    def __init__(self, power: int, repeat: bool = False, n: int = 1) -> None:
         super().__init__("recharge")
         self.power = power
         self.repeat = repeat
+        self.n = n
         self.power_cost = 0
 
     def state_dict(self):
-        return np.array([5, 0, 0, self.power, self.repeat])
+        return np.array([5, 0, 0, self.power, self.repeat, self.n])
 
 
 def format_factory_action(a: int):
@@ -138,17 +163,17 @@ def format_action_vec(a: np.ndarray):
     # (0 = move, 1 = transfer X amount of R, 2 = pickup X amount of R, 3 = dig, 4 = self destruct, 5 = recharge X)
     a_type = a[0]
     if a_type == 0:
-        return MoveAction(a[1], dist=1, repeat=a[4])
+        return MoveAction(a[1], dist=1, repeat=a[4], n=a[5])
     elif a_type == 1:
-        return TransferAction(a[1], a[2], a[3], repeat=a[4])
+        return TransferAction(a[1], a[2], a[3], repeat=a[4], n=a[5])
     elif a_type == 2:
-        return PickupAction(a[2], a[3], repeat=a[4])
+        return PickupAction(a[2], a[3], repeat=a[4], n=a[5])
     elif a_type == 3:
-        return DigAction(repeat=a[4])
+        return DigAction(repeat=a[4], n=a[5])
     elif a_type == 4:
-        return SelfDestructAction(repeat=a[4])
+        return SelfDestructAction(repeat=a[4], n=a[5])
     elif a_type == 5:
-        return RechargeAction(a[3], repeat=a[4])
+        return RechargeAction(a[3], repeat=a[4], n=a[5])
     else:
         raise ValueError(f"Action {a} is invalid type, {a[0]} is not valid")
 
@@ -157,9 +182,10 @@ def format_action_vec(a: np.ndarray):
 move_deltas = np.array([[0, 0], [0, -1], [1, 0], [0, 1], [-1, 0]])
 
 
-def validate_actions(env_cfg: EnvConfig, state: 'State', actions_by_type, weather_cfg, verbose=1):
+def validate_actions(env_cfg: EnvConfig, state: "State", actions_by_type, verbose=1):
     """
-    validates actions and logs warnings for any invalid actions. Invalid actions are subsequently not evaluated
+    validates actions and logs warnings for any invalid actions. Invalid actions are subsequently not evaluated and
+    won't count towards the `n` value of actions that execute multiple times.
     """
     actions_by_type_validated = defaultdict(list)
     valid_action = True
@@ -167,7 +193,8 @@ def validate_actions(env_cfg: EnvConfig, state: 'State', actions_by_type, weathe
     def invalidate_action(msg):
         nonlocal valid_action
         valid_action = False
-        if verbose > 0: print(f"{state.real_env_steps}: {msg}")
+        if verbose > 0:
+            print(f"{state.real_env_steps}: {msg}")
 
     for unit, transfer_action in actions_by_type["transfer"]:
         valid_action = True
@@ -199,10 +226,10 @@ def invalidate_action(msg):
     for unit, dig_action in actions_by_type["dig"]:
         valid_action = True
         dig_action: DigAction
-        dig_cost = math.ceil(unit.unit_cfg.DIG_COST * weather_cfg["power_loss_factor"])
+        dig_cost = unit.unit_cfg.DIG_COST
         if dig_cost > unit.power:
             invalidate_action(
-                f"Invalid Dig Action for unit {unit} - Tried to dig requiring ceil({unit.unit_cfg.DIG_COST} x {weather_cfg['power_loss_factor']}) = {dig_cost} power but only had {unit.power} power. Power cost factor is {weather_cfg['power_loss_factor']}"
+                f"Invalid Dig Action for unit {unit} - Tried to dig requiring {dig_cost} power"
             )
             continue
         # verify not digging over a factory which is not allowed
@@ -228,7 +255,9 @@ def invalidate_action(msg):
     for unit, move_action in actions_by_type["move"]:
         valid_action = True
         move_action: MoveAction
-        target_pos: Position = unit.pos + move_action.dist * move_deltas[move_action.move_dir]
+        target_pos: Position = (
+            unit.pos + move_action.dist * move_deltas[move_action.move_dir]
+        )
         if (
             target_pos.x < 0
             or target_pos.y < 0
@@ -249,16 +278,12 @@ def invalidate_action(msg):
                 continue
         rubble = state.board.rubble[target_pos.x, target_pos.y]
         power_required = (
-            0
-            if move_action.move_dir == 0
-            else math.ceil(
-                unit.move_power_cost(rubble) * weather_cfg["power_loss_factor"]
-            )
+            0 if move_action.move_dir == 0 else unit.move_power_cost(rubble)
         )
 
         if power_required > unit.power:
             invalidate_action(
-                f"Invalid movement action for unit {unit} - Tried to move to {target_pos} requiring ceil({unit.move_power_cost(rubble)} x {weather_cfg['power_loss_factor']}) power but only had {unit.power} power. Power cost factor is {weather_cfg['power_loss_factor']}"
+                f"Invalid movement action for unit {unit} - Tried to move to {target_pos} requiring {power_required} power"
             )
             continue
         if valid_action:
@@ -268,15 +293,17 @@ def invalidate_action(msg):
     for unit, self_destruct_action in actions_by_type["self_destruct"]:
         valid_action = True
         self_destruct_action: SelfDestructAction
-        power_required = math.ceil(unit.unit_cfg.SELF_DESTRUCT_COST * weather_cfg["power_loss_factor"])
+        power_required = unit.unit_cfg.SELF_DESTRUCT_COST
         if power_required > unit.power:
             invalidate_action(
-                f"Invalid self destruct action for unit {unit} - Tried to self destruct requiring ceil({unit.unit_cfg.SELF_DESTRUCT_COST} x {weather_cfg['power_loss_factor']}) power but only had {unit.power} power. Power cost factor is {weather_cfg['power_loss_factor']}"
+                f"Invalid self destruct action for unit {unit} - Tried to self destruct requiring {power_required} power"
             )
             continue
         if valid_action:
             self_destruct_action.power_cost = power_required
-            actions_by_type_validated["self_destruct"].append((unit, self_destruct_action))
+            actions_by_type_validated["self_destruct"].append(
+                (unit, self_destruct_action)
+            )
 
     for unit, recharge_action in actions_by_type["recharge"]:
         actions_by_type_validated["recharge"].append((unit, recharge_action))
@@ -284,15 +311,19 @@ def invalidate_action(msg):
     for factory, build_action in actions_by_type["factory_build"]:
         valid_action = True
         build_action: FactoryBuildAction
-        factory: 'Factory'
+        factory: "Factory"
 
         unit_cfg = env_cfg.ROBOTS[build_action.unit_type.name]
         if factory.cargo.metal < unit_cfg.METAL_COST:
-            invalidate_action(f"Invalid factory build action for factory {factory} - Insufficient metal, factory has {factory.cargo.metal}, but requires {unit_cfg.METAL_COST} to build {build_action.unit_type}")
+            invalidate_action(
+                f"Invalid factory build action for factory {factory} - Insufficient metal, factory has {factory.cargo.metal}, but requires {unit_cfg.METAL_COST} to build {build_action.unit_type}"
+            )
             continue
-        power_required = math.ceil(unit_cfg.POWER_COST * weather_cfg["power_loss_factor"])
+        power_required = unit_cfg.POWER_COST
         if factory.power < power_required:
-            invalidate_action(f"Invalid factory build action for factory {factory} - Insufficient power, factory has {factory.power}, but requires ceil({unit_cfg.POWER_COST} x {weather_cfg['power_loss_factor']}) to build {build_action.unit_type.name}. Power cost factor is {weather_cfg['power_loss_factor']}")
+            invalidate_action(
+                f"Invalid factory build action for factory {factory} - Insufficient power, factory has {factory.power}, but requires {power_required} power"
+            )
             continue
         if valid_action:
             build_action.power_cost = power_required
@@ -303,7 +334,9 @@ def invalidate_action(msg):
         valid_action = True
         water_cost = factory.water_cost(env_cfg)
         if water_cost > factory.cargo.water:
-            invalidate_action(f"Invalid factory water action for factory {factory} - Insufficient water, factory has {factory.cargo.water}, but requires {water_cost} to water lichen")
+            invalidate_action(
+                f"Invalid factory water action for factory {factory} - Insufficient water, factory has {factory.cargo.water}, but requires {water_cost} to water lichen"
+            )
             continue
         if valid_action:
             actions_by_type_validated["factory_water"].append((factory, water_action))
diff --git a/luxai_s2/luxai_s2/config.py b/luxai_s2/luxai_s2/config.py
index 97cf6b61..1ea9ef1b 100644
--- a/luxai_s2/luxai_s2/config.py
+++ b/luxai_s2/luxai_s2/config.py
@@ -1,6 +1,6 @@
+import dataclasses
 from argparse import Namespace
 from dataclasses import dataclass
-import dataclasses
 from typing import Dict, List
 
 
@@ -20,7 +20,7 @@ class UnitConfig:
     CHARGE: int = 1
     INIT_POWER: int = 50
     MOVE_COST: int = 1
-    RUBBLE_MOVEMENT_COST: int = 1
+    RUBBLE_MOVEMENT_COST: float = 1
     DIG_COST: int = 5
     DIG_RUBBLE_REMOVED: int = 1
     DIG_RESOURCE_GAIN: int = 2
@@ -30,17 +30,15 @@ class UnitConfig:
     ACTION_QUEUE_POWER_COST: int = 1
 
 
-
 @dataclass
 class EnvConfig:
     ## various options that can be configured if needed
 
-
     ### Variable parameters that don't affect game logic much ###
     max_episode_length: int = 1000
     map_size: int = 48
     verbose: int = 1
-    
+
     # this can be disabled to improve env FPS but assume your actions are well formatted
     # During online competition this is set to True
     validate_action_space: bool = True
@@ -52,19 +50,22 @@ class EnvConfig:
     MAX_FACTORIES: int = 5
     CYCLE_LENGTH: int = 50
     DAY_LENGTH: int = 30
-    UNIT_ACTION_QUEUE_SIZE: int = 20 # when set to 1, then no action queue is used
-
+    UNIT_ACTION_QUEUE_SIZE: int = 20  # when set to 1, then no action queue is used
 
     MAX_RUBBLE: int = 100
     FACTORY_RUBBLE_AFTER_DESTRUCTION: int = 50
-    INIT_WATER_METAL_PER_FACTORY: int = 150 # amount of water and metal units given to each factory
+    INIT_WATER_METAL_PER_FACTORY: int = (
+        150  # amount of water and metal units given to each factory
+    )
     INIT_POWER_PER_FACTORY: int = 1000
 
     #### LICHEN ####
-    MIN_LICHEN_TO_SPREAD: int = 10
+    MIN_LICHEN_TO_SPREAD: int = 20
     LICHEN_LOST_WITHOUT_WATER: int = 1
     LICHEN_GAINED_WITH_WATER: int = 1
     MAX_LICHEN_PER_TILE: int = 100
+    POWER_PER_CONNECTED_LICHEN_TILE: int = 1
+
     # cost of watering with a factory is `ceil(# of connected lichen tiles) / (this factor) + 1`
     LICHEN_WATERING_COST_FACTOR: int = 10
 
@@ -72,75 +73,64 @@ class EnvConfig:
     BIDDING_SYSTEM: bool = True
 
     #### Factories ####
-    FACTORY_PROCESSING_RATE_WATER: int = 50
-    ICE_WATER_RATIO: int = 2
+    FACTORY_PROCESSING_RATE_WATER: int = 100
+    ICE_WATER_RATIO: int = 4
     FACTORY_PROCESSING_RATE_METAL: int = 50
     ORE_METAL_RATIO: int = 5
     # game design note: Factories close to resource cluster = more resources are refined per turn
-    # Then the high ice:water and ore:metal ratios encourages transfer of refined resources between 
+    # Then the high ice:water and ore:metal ratios encourages transfer of refined resources between
     # factories dedicated to mining particular clusters which is more possible as it is more compact
-    # 
-
 
     FACTORY_CHARGE: int = 50
     FACTORY_WATER_CONSUMPTION: int = 1
     # game design note: with a positve water consumption, game becomes quite hard for new competitors.
     # so we set it to 0
 
+    #### Collision Mechanics ####
+    POWER_LOSS_FACTOR: float = 0.5
 
     #### Units ####
     ROBOTS: Dict[str, UnitConfig] = dataclasses.field(
         default_factory=lambda: dict(
-        LIGHT=UnitConfig(
-            METAL_COST=10, POWER_COST=50, INIT_POWER=50, CARGO_SPACE=100, BATTERY_CAPACITY=150, CHARGE=1, MOVE_COST=1, RUBBLE_MOVEMENT_COST=0,
-            DIG_COST=5,
-            SELF_DESTRUCT_COST=5,
-            DIG_RUBBLE_REMOVED=2,
-            DIG_RESOURCE_GAIN=2,
-            DIG_LICHEN_REMOVED=10,
-            RUBBLE_AFTER_DESTRUCTION=1,
-            ACTION_QUEUE_POWER_COST=1,
-        ),
-
-        HEAVY=UnitConfig(
-            METAL_COST=100, POWER_COST=500, INIT_POWER=500, CARGO_SPACE=1000, BATTERY_CAPACITY=3000, CHARGE=10, MOVE_COST=20, RUBBLE_MOVEMENT_COST=1,
-            DIG_COST=100,
-            SELF_DESTRUCT_COST=100,
-            DIG_RUBBLE_REMOVED=20,
-            DIG_RESOURCE_GAIN=20,
-            DIG_LICHEN_REMOVED=100,
-            RUBBLE_AFTER_DESTRUCTION=10,
-            ACTION_QUEUE_POWER_COST=10,
-        ),
+            LIGHT=UnitConfig(
+                METAL_COST=10,
+                POWER_COST=50,
+                INIT_POWER=50,
+                CARGO_SPACE=100,
+                BATTERY_CAPACITY=150,
+                CHARGE=1,
+                MOVE_COST=1,
+                RUBBLE_MOVEMENT_COST=0.05,
+                DIG_COST=5,
+                SELF_DESTRUCT_COST=5,
+                DIG_RUBBLE_REMOVED=2,
+                DIG_RESOURCE_GAIN=2,
+                DIG_LICHEN_REMOVED=10,
+                RUBBLE_AFTER_DESTRUCTION=1,
+                ACTION_QUEUE_POWER_COST=1,
+            ),
+            HEAVY=UnitConfig(
+                METAL_COST=100,
+                POWER_COST=500,
+                INIT_POWER=500,
+                CARGO_SPACE=1000,
+                BATTERY_CAPACITY=3000,
+                CHARGE=10,
+                MOVE_COST=20,
+                RUBBLE_MOVEMENT_COST=1,
+                DIG_COST=60,
+                SELF_DESTRUCT_COST=100,
+                DIG_RUBBLE_REMOVED=20,
+                DIG_RESOURCE_GAIN=20,
+                DIG_LICHEN_REMOVED=100,
+                RUBBLE_AFTER_DESTRUCTION=10,
+                ACTION_QUEUE_POWER_COST=10,
+            ),
         )
     )
 
-    #### Weather ####
-    WEATHER_ID_TO_NAME: list = dataclasses.field(default_factory=lambda : ["NONE","MARS_QUAKE","COLD_SNAP","DUST_STORM","SOLAR_FLARE"])
-    NUM_WEATHER_EVENTS_RANGE: List[int] = dataclasses.field(default_factory=lambda:[3,5])
-    WEATHER: dict = dataclasses.field(default_factory=lambda:dict(
-        MARS_QUAKE=dict(
-            TIME_RANGE=[1, 5]
-        ),
-        COLD_SNAP=dict(
-            # power multiplier required per robot action. 2 -> requires 2x as much power to execute the same action
-            POWER_CONSUMPTION=2,
-            TIME_RANGE=[10, 30]
-        ),
-        DUST_STORM=dict(
-            # power gain multiplier. .5 -> gain .5x as much power per turn
-            POWER_GAIN=0.5,
-            TIME_RANGE=[10, 30]
-        ),
-        SOLAR_FLARE=dict(
-            # power gain multiplier. 2 -> gain 2x as much power per turn
-            POWER_GAIN=2,
-            TIME_RANGE=[10, 30]
-        ),
-    ))
-
     @classmethod
     def from_dict(cls, data):
         data["ROBOTS"]["LIGHT"] = UnitConfig(**data["ROBOTS"]["LIGHT"])
         data["ROBOTS"]["HEAVY"] = UnitConfig(**data["ROBOTS"]["HEAVY"])
-        return cls(**data)
\ No newline at end of file
+        return cls(**data)
diff --git a/luxai_s2/luxai_s2/env.py b/luxai_s2/luxai_s2/env.py
index 8d465511..a66af7cf 100644
--- a/luxai_s2/luxai_s2/env.py
+++ b/luxai_s2/luxai_s2/env.py
@@ -1,12 +1,13 @@
-from collections import OrderedDict, defaultdict
 import functools
 import math
-from typing import Dict, List, Set, Tuple, Union
 import traceback
+from collections import OrderedDict, defaultdict
+from typing import Dict, List, Set, Tuple, Union
+
 import numpy as np
 from pettingzoo import ParallelEnv
 from pettingzoo.utils import wrappers
-from luxai_s2.map import weather
+
 from luxai_s2.actions import (
     Action,
     DigAction,
@@ -19,25 +20,30 @@
     TransferAction,
     format_action_vec,
     format_factory_action,
-    validate_actions,
     move_deltas,
+    validate_actions,
 )
-
 from luxai_s2.config import EnvConfig
 from luxai_s2.factory import Factory
 from luxai_s2.map.board import Board
 from luxai_s2.map.position import Position
 from luxai_s2.pyvisual.visualizer import Visualizer
-from luxai_s2.spaces.act_space import get_act_space, get_act_space_bid, get_act_space_init, get_act_space_placement
+from luxai_s2.spaces.act_space import (
+    get_act_space,
+    get_act_space_bid,
+    get_act_space_init,
+    get_act_space_placement,
+)
 from luxai_s2.spaces.obs_space import get_obs_space
 from luxai_s2.state import State
 from luxai_s2.team import FactionTypes, Team
 from luxai_s2.unit import Unit, UnitType
-from luxai_s2.utils.utils import is_day
+from luxai_s2.utils.utils import get_top_two_power_units, is_day
 
 # some utility types
 ActionsByType = Dict[str, List[Tuple[Unit, Action]]]
 
+
 def env():
     """
     The env function often wraps the environment in wrappers by default.
@@ -59,15 +65,19 @@ class LuxAI_S2(ParallelEnv):
     metadata = {"render.modes": ["human", "html", "rgb_array"], "name": "luxai_s2_v0"}
 
     def __init__(self, collect_stats: bool = False, **kwargs):
-        self.collect_stats = collect_stats # note: added here instead of in configs since it would break existing bots
+        self.collect_stats = collect_stats  # note: added here instead of in configs since it would break existing bots
         default_config = EnvConfig(**kwargs)
 
         self.env_cfg = default_config
         self.possible_agents = ["player_" + str(r) for r in range(2)]
-        self.agent_name_mapping = dict(zip(self.possible_agents, list(range(len(self.possible_agents)))))
+        self.agent_name_mapping = dict(
+            zip(self.possible_agents, list(range(len(self.possible_agents))))
+        )
         self.max_episode_length = self.env_cfg.max_episode_length
 
-        self.state: State = State(seed_rng=None, seed=-1, env_cfg=self.env_cfg, env_steps=-1, board=None, weather_schedule=None)
+        self.state: State = State(
+            seed_rng=None, seed=-1, env_cfg=self.env_cfg, env_steps=-1, board=None
+        )
 
         self.seed_rng: np.random.RandomState = None
 
@@ -77,7 +87,9 @@ def __init__(self, collect_stats: bool = False, **kwargs):
     # allows action space seeding to work as expected
     @functools.lru_cache(maxsize=None)
     def observation_space(self, agent):
-        return get_obs_space(config=self.env_cfg, agent_names=self.possible_agents, agent=agent)
+        return get_obs_space(
+            config=self.env_cfg, agent_names=self.possible_agents, agent=agent
+        )
 
     # @functools.lru_cache(maxsize=None)
     def action_space(self, agent: str):
@@ -87,17 +99,22 @@ def action_space(self, agent: str):
                 return get_act_space_bid(config=self.env_cfg, agent=agent)
             if self.env_steps <= self.state.board.factories_per_team + 1:
                 return get_act_space_placement(config=self.env_cfg, agent=agent)
-            return get_act_space(self.state.units, self.state.factories, config=self.env_cfg, agent=agent)
+            return get_act_space(
+                self.state.units, self.state.factories, config=self.env_cfg, agent=agent
+            )
         else:
             if self.env_steps == 0:
                 return get_act_space_init(config=self.env_cfg, agent=agent)
-            return get_act_space(self.state.units, self.state.factories, config=self.env_cfg, agent=agent)
+            return get_act_space(
+                self.state.units, self.state.factories, config=self.env_cfg, agent=agent
+            )
 
     def _init_render(self):
         if self.py_visualizer is None:
             self.py_visualizer = Visualizer(self.state)
             return True
         return False
+
     def render(self, mode="human", **kwargs):
         """
         Renders the environment. In human mode, it can print to terminal, open
@@ -107,7 +124,7 @@ def render(self, mode="human", **kwargs):
         if mode == "human":
             if self._init_render():
                 self.py_visualizer.init_window()
-            
+
             self.py_visualizer.update_scene(self.state)
             self.py_visualizer.render()
         elif mode == "rgb_array":
@@ -119,7 +136,9 @@ def render(self, mode="human", **kwargs):
                 VIDEO_W = kwargs["width"]
             if "height" in kwargs:
                 VIDEO_H = kwargs["height"]
-            return self.py_visualizer._create_image_array(self.py_visualizer.surf, (VIDEO_W, VIDEO_H))
+            return self.py_visualizer._create_image_array(
+                self.py_visualizer.surf, (VIDEO_W, VIDEO_H)
+            )
 
     def close(self):
         """
@@ -129,6 +148,7 @@ def close(self):
         """
         try:
             import pygame
+
             pygame.display.quit()
             pygame.quit()
         except:
@@ -162,17 +182,24 @@ def reset(self, seed=None):
         self.env_steps = 0
         self.seed = seed
         board = Board(seed=seed, env_cfg=self.env_cfg)
-        weather_schedule = weather.generate_weather_schedule(seed_rng, self.state.env_cfg)
-        self.state: State = State(seed_rng=seed_rng, seed=seed, env_cfg=self.state.env_cfg, env_steps=0, board=board, weather_schedule=weather_schedule)
+        self.state: State = State(
+            seed_rng=seed_rng,
+            seed=seed,
+            env_cfg=self.state.env_cfg,
+            env_steps=0,
+            board=board,
+        )
         self.max_episode_length = self.env_cfg.max_episode_length
         for agent in self.possible_agents:
             self.state.units[agent] = OrderedDict()
             self.state.factories[agent] = OrderedDict()
             if self.collect_stats:
                 self.state.stats[agent] = dict(
-                    lichen_grown=0, lichen_lost=0, 
-                    units_lost=dict(LIGHT=0, HEAVY=0), 
-                    factories_lost=0, units_built=dict(LIGHT=0, HEAVY=0)
+                    lichen_grown=0,
+                    lichen_lost=0,
+                    units_lost=dict(LIGHT=0, HEAVY=0),
+                    factories_lost=0,
+                    units_built=dict(LIGHT=0, HEAVY=0),
                 )
         obs = self.state.get_obs()
         observations = {agent: obs for agent in self.agents}
@@ -194,15 +221,29 @@ def _handle_bid(self, actions):
                 raise ValueError(f"Invalid player {k}")
             if "faction" in a and "bid" in a:
                 if a["faction"] not in [e.name for e in FactionTypes]:
-                    self._log(f"{k} initialized with invalid faction name {a['faction']}")
+                    self._log(
+                        f"{k} initialized with invalid faction name {a['faction']}"
+                    )
                     failed_agents[k] = True
                     continue
                 self.state.teams[k] = Team(
-                    team_id=self.agent_name_mapping[k], agent=k, faction=FactionTypes[a["faction"]]
+                    team_id=self.agent_name_mapping[k],
+                    agent=k,
+                    faction=FactionTypes[a["faction"]],
                 )
-                self.state.teams[k].init_water = self.env_cfg.INIT_WATER_METAL_PER_FACTORY * (self.state.board.factories_per_team)
-                self.state.teams[k].init_metal = self.env_cfg.INIT_WATER_METAL_PER_FACTORY * (self.state.board.factories_per_team)
-                self.state.teams[k].factories_to_place = self.state.board.factories_per_team
+                self.state.teams[
+                    k
+                ].init_water = self.env_cfg.INIT_WATER_METAL_PER_FACTORY * (
+                    self.state.board.factories_per_team
+                )
+                self.state.teams[
+                    k
+                ].init_metal = self.env_cfg.INIT_WATER_METAL_PER_FACTORY * (
+                    self.state.board.factories_per_team
+                )
+                self.state.teams[
+                    k
+                ].factories_to_place = self.state.board.factories_per_team
                 # verify bid is valid
                 valid_action = True
                 bid = abs(a["bid"])
@@ -243,6 +284,7 @@ def _handle_bid(self, actions):
                 self.state.teams[highest_bid_agent].place_first = True
                 self.state.teams[lowest_bid_agent].place_first = False
         return failed_agents
+
     def _handle_factory_placement_step(self, actions):
         # factory placement rounds, which are sequential
 
@@ -258,7 +300,6 @@ def _handle_factory_placement_step(self, actions):
             else:
                 player_to_place_factory = "player_0"
 
-
         failed_agents = {agent: False for agent in self.agents}
         for k, a in actions.items():
             if a is None:
@@ -268,22 +309,33 @@ def _handle_factory_placement_step(self, actions):
                 raise ValueError(f"Invalid player {k}")
             if "spawn" in a and "metal" in a and "water" in a:
                 if k != player_to_place_factory:
-                    self._log(f"{k} tried to perform an action in the early phase when it is not its turn right now.")
+                    self._log(
+                        f"{k} tried to perform an action in the early phase when it is not its turn right now."
+                    )
                     continue
                 if self.state.teams[k].factories_to_place <= 0:
-                    self._log(f"{k} cannot place additional factories. Cancelled placement of factory")
+                    self._log(
+                        f"{k} cannot place additional factories. Cancelled placement of factory"
+                    )
                     continue
                 if a["water"] < 0 or a["metal"] < 0:
-                    self._log(f"{k} tried to place negative water/metal in factory. Cancelled placement of factory")
+                    self._log(
+                        f"{k} tried to place negative water/metal in factory. Cancelled placement of factory"
+                    )
                     continue
                 if a["water"] > self.state.teams[k].init_water:
                     a["water"] = self.state.teams[k].init_water
-                    self._log(f" Warning - {k} does not have enough water. Using {a['water']}")
+                    self._log(
+                        f" Warning - {k} does not have enough water. Using {a['water']}"
+                    )
                 if a["metal"] > self.state.teams[k].init_metal:
                     a["metal"] = self.state.teams[k].init_metal
-                    self._log(f" Warning - {k} does not have enough metal. Using {a['metal']}")
+                    self._log(
+                        f" Warning - {k} does not have enough metal. Using {a['metal']}"
+                    )
                 factory = self.add_factory(self.state.teams[k], a["spawn"])
-                if factory is None: continue
+                if factory is None:
+                    continue
                 factory.cargo.water = a["water"]
                 factory.cargo.metal = a["metal"]
                 factory.power = self.env_cfg.INIT_POWER_PER_FACTORY
@@ -294,6 +346,7 @@ def _handle_factory_placement_step(self, actions):
                 # pass, turn is skipped.
                 pass
         return failed_agents
+
     def _handle_nobidding_early_game(self, actions):
         failed_agents = {agent: False for agent in self.agents}
         for k, a in actions.items():
@@ -304,15 +357,21 @@ def _handle_nobidding_early_game(self, actions):
                 raise ValueError(f"Invalid player {k}")
             if "spawns" in a and "faction" in a:
                 self.state.teams[k] = Team(
-                    team_id=self.agent_name_mapping[k], agent=k, faction=FactionTypes[a["faction"]]
+                    team_id=self.agent_name_mapping[k],
+                    agent=k,
+                    faction=FactionTypes[a["faction"]],
                 )
                 if len(a["spawns"]) > self.state.board.factories_per_team:
-                    if self.env_cfg.verbose > 0: self._log(f"{k} tried to spawn more factories than allocated in board.factories_per_team. Spawning only the first {self.state.board.factories_per_team} locations")
-                for spawn_loc in a["spawns"][:self.state.board.factories_per_team]:
+                    if self.env_cfg.verbose > 0:
+                        self._log(
+                            f"{k} tried to spawn more factories than allocated in board.factories_per_team. Spawning only the first {self.state.board.factories_per_team} locations"
+                        )
+                for spawn_loc in a["spawns"][: self.state.board.factories_per_team]:
                     self.add_factory(self.state.teams[k], spawn_loc)
             else:
                 # team k loses
                 failed_agents[k] = True
+
     def _step_early_game(self, actions):
         # handle initialization
         failed_agents = {agent: False for agent in self.agents}
@@ -334,49 +393,75 @@ def _handle_transfer_actions(self, actions_by_type: ActionsByType):
         amount_list = []
         for unit, transfer_action in actions_by_type["transfer"]:
             transfer_action: TransferAction
-            transfer_amount = unit.sub_resource(transfer_action.resource, transfer_action.transfer_amount)
+            transfer_amount = unit.sub_resource(
+                transfer_action.resource, transfer_action.transfer_amount
+            )
             amount_list.append(transfer_amount)
 
         # add to target cargo
-        for (unit, transfer_action), transfer_amount in zip(actions_by_type["transfer"], amount_list):
+        for (unit, transfer_action), transfer_amount in zip(
+            actions_by_type["transfer"], amount_list
+        ):
             transfer_action: TransferAction
-            transfer_pos: Position = unit.pos + move_deltas[transfer_action.transfer_dir]
+            transfer_pos: Position = (
+                unit.pos + move_deltas[transfer_action.transfer_dir]
+            )
             units_there = self.state.board.get_units_at(transfer_pos)
 
             # if there is a factory, we prefer transferring to that entity
             factory_id = f"factory_{self.state.board.factory_occupancy_map[transfer_pos.x, transfer_pos.y]}"
             if factory_id in self.state.factories[unit.team.agent]:
                 factory = self.state.factories[unit.team.agent][factory_id]
-                actually_transferred = factory.add_resource(transfer_action.resource, transfer_amount)
+                actually_transferred = factory.add_resource(
+                    transfer_action.resource, transfer_amount
+                )
             elif units_there is not None:
                 assert len(units_there) == 1, "Fatal error here, this is a bug"
                 target_unit = units_there[0]
                 # add resources to target. This will waste (transfer_amount - actually_transferred) resources
-                actually_transferred = target_unit.add_resource(transfer_action.resource, transfer_amount)
+                actually_transferred = target_unit.add_resource(
+                    transfer_action.resource, transfer_amount
+                )
             unit.repeat_action(transfer_action)
 
     def _handle_pickup_actions(self, actions_by_type: ActionsByType):
         for unit, pickup_action in actions_by_type["pickup"]:
             pickup_action: PickupAction
             factory = self.state.board.get_factory_at(self.state, unit.pos)
-            pickup_amount = factory.sub_resource(pickup_action.resource, pickup_action.pickup_amount)
+            pickup_amount = factory.sub_resource(
+                pickup_action.resource, pickup_action.pickup_amount
+            )
             # may waste resources if tried to pickup more than one can hold.
             actually_pickedup = unit.add_resource(pickup_action.resource, pickup_amount)
             unit.repeat_action(pickup_action)
 
-    def _handle_dig_actions(self, actions_by_type: ActionsByType, weather_cfg):
+    def _handle_dig_actions(self, actions_by_type: ActionsByType):
         for unit, dig_action in actions_by_type["dig"]:
             dig_action: DigAction
             if self.state.board.rubble[unit.pos.x, unit.pos.y] > 0:
-                self.state.board.rubble[unit.pos.x, unit.pos.y] = max(self.state.board.rubble[unit.pos.x, unit.pos.y] - unit.unit_cfg.DIG_RUBBLE_REMOVED, 0)
+                self.state.board.rubble[unit.pos.x, unit.pos.y] = max(
+                    self.state.board.rubble[unit.pos.x, unit.pos.y]
+                    - unit.unit_cfg.DIG_RUBBLE_REMOVED,
+                    0,
+                )
             elif self.state.board.lichen[unit.pos.x, unit.pos.y] > 0:
-                self.state.board.lichen[unit.pos.x, unit.pos.y] = max(self.state.board.lichen[unit.pos.x, unit.pos.y] - unit.unit_cfg.DIG_LICHEN_REMOVED, 0)
+                lichen_left = max(
+                    self.state.board.lichen[unit.pos.x, unit.pos.y]
+                    - unit.unit_cfg.DIG_LICHEN_REMOVED,
+                    0,
+                )
+                self.state.board.lichen[unit.pos.x, unit.pos.y] = lichen_left
+                if lichen_left == 0:  # dug out the last lichen
+                    self.state.board.rubble[
+                        unit.pos.x, unit.pos.y
+                    ] = self.state.env_cfg.ROBOTS[unit.unit_type.name].DIG_RESOURCE_GAIN
             elif self.state.board.ice[unit.pos.x, unit.pos.y] > 0:
                 unit.add_resource(0, unit.unit_cfg.DIG_RESOURCE_GAIN)
             elif self.state.board.ore[unit.pos.x, unit.pos.y] > 0:
                 unit.add_resource(1, unit.unit_cfg.DIG_RESOURCE_GAIN)
-            unit.power -= math.ceil(self.state.env_cfg.ROBOTS[unit.unit_type.name].DIG_COST * weather_cfg["power_loss_factor"])
+            unit.power -= self.state.env_cfg.ROBOTS[unit.unit_type.name].DIG_COST
             unit.repeat_action(dig_action)
+
     def _handle_self_destruct_actions(self, actions_by_type: ActionsByType):
         for unit, self_destruct_action in actions_by_type["self_destruct"]:
             unit: Unit
@@ -385,9 +470,11 @@ def _handle_self_destruct_actions(self, actions_by_type: ActionsByType):
             del self.state.board.units_map[pos_hash]
             self.destroy_unit(unit)
             if self.collect_stats:
-                self.state.stats[unit.team.agent]["units_lost"][unit.unit_type.name] += 1
+                self.state.stats[unit.team.agent]["units_lost"][
+                    unit.unit_type.name
+                ] += 1
 
-    def _handle_factory_build_actions(self, actions_by_type: ActionsByType, weather_cfg):
+    def _handle_factory_build_actions(self, actions_by_type: ActionsByType):
         for factory, factory_build_action in actions_by_type["factory_build"]:
             factory: Factory
             factory_build_action: FactoryBuildAction
@@ -398,31 +485,42 @@ def _handle_factory_build_actions(self, actions_by_type: ActionsByType, weather_
 
             if factory_build_action.unit_type == UnitType.HEAVY:
                 if factory.cargo.metal < self.env_cfg.ROBOTS["HEAVY"].METAL_COST:
-                    self._log(f"{factory} doesn't have enough metal to build a heavy despite having enough metal at the start of the turn. This is likely because a unit picked up some of the metal.")
+                    self._log(
+                        f"{factory} doesn't have enough metal to build a heavy despite having enough metal at the start of the turn. This is likely because a unit picked up some of the metal."
+                    )
                     continue
                 if factory.power < factory_build_action.power_cost:
-                    self._log(f"{factory} doesn't have enough power to build a heavy despite having enough power at the start of the turn. This is likely because a unit picked up some of the power.")
+                    self._log(
+                        f"{factory} doesn't have enough power to build a heavy despite having enough power at the start of the turn. This is likely because a unit picked up some of the power."
+                    )
                     continue
                 factory.sub_resource(3, self.env_cfg.ROBOTS["HEAVY"].METAL_COST)
                 factory.sub_resource(4, factory_build_action.power_cost)
             else:
                 if factory.cargo.metal < self.env_cfg.ROBOTS["LIGHT"].METAL_COST:
-                    self._log(f"{factory} doesn't have enough metal to build a light despite having enough metal at the start of the turn. This is likely because a unit picked up some of the metal.")
+                    self._log(
+                        f"{factory} doesn't have enough metal to build a light despite having enough metal at the start of the turn. This is likely because a unit picked up some of the metal."
+                    )
                     continue
                 if factory.power < factory_build_action.power_cost:
-                    self._log(f"{factory} doesn't have enough power to build a light despite having enough power at the start of the turn. This is likely because a unit picked up some of the power.")
+                    self._log(
+                        f"{factory} doesn't have enough power to build a light despite having enough power at the start of the turn. This is likely because a unit picked up some of the power."
+                    )
                     continue
                 factory.sub_resource(3, self.env_cfg.ROBOTS["LIGHT"].METAL_COST)
                 factory.sub_resource(4, factory_build_action.power_cost)
-            
+
             self.add_unit(
                 team=team,
                 unit_type=factory_build_action.unit_type,
                 pos=factory.pos.pos,
             )
             if self.collect_stats:
-                self.state.stats[factory.team.agent]["units_built"][factory_build_action.unit_type.name] += 1
-    def _handle_movement_actions(self, actions_by_type: ActionsByType, weather_cfg):
+                self.state.stats[factory.team.agent]["units_built"][
+                    factory_build_action.unit_type.name
+                ] += 1
+
+    def _handle_movement_actions(self, actions_by_type: ActionsByType):
         new_units_map: Dict[str, List[Unit]] = defaultdict(list)
         heavy_entered_pos: Dict[str, List[Unit]] = defaultdict(list)
         light_entered_pos: Dict[str, List[Unit]] = defaultdict(list)
@@ -432,7 +530,9 @@ def _handle_movement_actions(self, actions_by_type: ActionsByType, weather_cfg):
             # skip move center
             if move_action.move_dir != 0:
                 old_pos_hash = self.state.board.pos_hash(unit.pos)
-                target_pos = unit.pos + move_action.dist * move_deltas[move_action.move_dir]
+                target_pos = (
+                    unit.pos + move_action.dist * move_deltas[move_action.move_dir]
+                )
                 power_required = move_action.power_cost
                 unit.pos = target_pos
                 new_pos_hash = self.state.board.pos_hash(unit.pos)
@@ -465,10 +565,27 @@ def _handle_movement_actions(self, actions_by_type: ActionsByType, weather_cfg):
                 new_units_map_after_collision[pos_hash] += units
                 continue
             if len(heavy_entered_pos[pos_hash]) > 1:
-                # all units collide and break
-                for u in units:
-                    destroyed_units.add(u)
-                self._log(f"{len(destroyed_units)} Units collided at {pos_hash}")
+                # all units collide, find the top 2 units by power
+                (most_power_unit, next_most_power_unit) = get_top_two_power_units(units)
+                if most_power_unit.power == next_most_power_unit.power:
+                    # tie, all units break
+                    for u in units:
+                        destroyed_units.add(u)
+                    self._log(
+                        f"{len(destroyed_units)} Units: ({', '.join([u.unit_id for u in destroyed_units])}) collided at {pos_hash}"
+                    )
+                else:
+                    most_power_unit_power_loss = math.ceil(
+                        next_most_power_unit.power * self.env_cfg.POWER_LOSS_FACTOR
+                    )
+                    most_power_unit.power -= most_power_unit_power_loss
+                    surviving_unit = most_power_unit
+                    for u in units:
+                        if u.unit_id != surviving_unit.unit_id:
+                            destroyed_units.add(u)
+                    self._log(
+                        f"{len(destroyed_units)} Units: ({', '.join([u.unit_id for u in destroyed_units])}) collided at {pos_hash} with {surviving_unit} surviving with {surviving_unit.power} power"
+                    )
                 all_destroyed_units.update(destroyed_units)
             elif len(heavy_entered_pos[pos_hash]) > 0:
                 # all other units collide and break
@@ -476,7 +593,9 @@ def _handle_movement_actions(self, actions_by_type: ActionsByType, weather_cfg):
                 for u in units:
                     if u.unit_id != surviving_unit.unit_id:
                         destroyed_units.add(u)
-                self._log(f"{len(destroyed_units)} Units: ({', '.join([u.unit_id for u in destroyed_units])}) collided at {pos_hash} with {surviving_unit} surviving")
+                self._log(
+                    f"{len(destroyed_units)} Units: ({', '.join([u.unit_id for u in destroyed_units])}) collided at {pos_hash} with {surviving_unit} surviving with {surviving_unit.power} power"
+                )
                 new_units_map_after_collision[pos_hash].append(surviving_unit)
                 all_destroyed_units.update(destroyed_units)
             else:
@@ -488,7 +607,10 @@ def _handle_movement_actions(self, actions_by_type: ActionsByType, weather_cfg):
                         if heavy_stationary_unit is not None:
                             heavy_stationary_unit = None
                             # we found >= 2 heavies stationary in a tile where no heavies are entering.
-                            self._log(f"At {pos_hash}, >= 2 heavies crashed as they were all stationary")
+                            # should only happen when spawning units
+                            self._log(
+                                f"At {pos_hash}, >= 2 heavies crashed as they were all stationary"
+                            )
                             break
                         heavy_stationary_unit = u
 
@@ -496,21 +618,47 @@ def _handle_movement_actions(self, actions_by_type: ActionsByType, weather_cfg):
                     surviving_unit = heavy_stationary_unit
                 else:
                     if len(light_entered_pos[pos_hash]) > 1:
-                        # all units collide
-                        surviving_unit = None
+                        # all units collide, get top 2 units by power
+                        (
+                            most_power_unit,
+                            next_most_power_unit,
+                        ) = get_top_two_power_units(units)
+                        print(
+                            most_power_unit.unit_id,
+                            most_power_unit.power,
+                            next_most_power_unit.unit_id,
+                            next_most_power_unit.power,
+                        )
+                        if most_power_unit.power == next_most_power_unit.power:
+                            # tie, all units break
+                            for u in units:
+                                destroyed_units.add(u)
+                        else:
+                            most_power_unit_power_loss = math.ceil(
+                                next_most_power_unit.power
+                                * self.env_cfg.POWER_LOSS_FACTOR
+                            )
+                            most_power_unit.power -= most_power_unit_power_loss
+                            surviving_unit = most_power_unit
                     elif len(light_entered_pos[pos_hash]) > 0:
                         # light crashes into stationary light unit
                         surviving_unit = light_entered_pos[pos_hash][0]
                 if surviving_unit is None:
                     for u in units:
+                        print(u.unit_id, u.power)
                         destroyed_units.add(u)
-                    self._log(f"{len(destroyed_units)} Units collided at {pos_hash}")
+                    self._log(
+                        f"{len(destroyed_units)} Units: ({', '.join([u.unit_id for u in destroyed_units])}) collided at {pos_hash}"
+                    )
                     all_destroyed_units.update(destroyed_units)
                 else:
                     for u in units:
                         if u.unit_id != surviving_unit.unit_id:
+                            print(u.unit_id, u.power)
                             destroyed_units.add(u)
-                    self._log(f"{len(destroyed_units)} Units collided at {pos_hash} with {surviving_unit} surviving")
+                    self._log(
+                        f"{len(destroyed_units)} Units: ({', '.join([u.unit_id for u in destroyed_units])}) collided at {pos_hash} with {surviving_unit} surviving with {surviving_unit.power} power"
+                    )
                     new_units_map_after_collision[pos_hash].append(surviving_unit)
                     all_destroyed_units.update(destroyed_units)
         self.state.board.units_map = new_units_map_after_collision
@@ -529,17 +677,24 @@ def _handle_recharge_actions(self, actions_by_type: ActionsByType):
             else:
                 # if unit got enough power, handle the action and consider it for repeating
                 unit.repeat_action(recharge_action)
+
     def _handle_factory_water_actions(self, actions_by_type: ActionsByType):
         for factory, factory_water_action in actions_by_type["factory_water"]:
             factory_water_action: FactoryWaterAction
             water_cost = factory.water_cost(self.env_cfg)
             if water_cost > factory.cargo.water:
-                self._log(f"{factory} has insufficient water to grow lichen, factory has {factory.cargo.water}, but requires {water_cost} to water lichen. This cost may have changed a little during this turn due to rubble changes and new tiles being grown on")
+                self._log(
+                    f"{factory} has insufficient water to grow lichen, factory has {factory.cargo.water}, but requires {water_cost} to water lichen. This cost may have changed a little during this turn due to rubble changes and new tiles being grown on"
+                )
                 continue
-            factory.cargo.water -= water_cost 
-            indexable_positions = ([v[0] for v in factory.grow_lichen_positions], [v[1] for v in factory.grow_lichen_positions])
+            factory.cargo.water -= water_cost
+            indexable_positions = (
+                [v[0] for v in factory.grow_lichen_positions],
+                [v[1] for v in factory.grow_lichen_positions],
+            )
             self.state.board.lichen[indexable_positions] += 2
             self.state.board.lichen_strains[indexable_positions] = factory.num_id
+
     def step(self, actions):
         """
         step(action) takes in an action for each agent and should return the
@@ -565,53 +720,68 @@ def step(self, actions):
         if early_game:
             failed_agents = self._step_early_game(actions)
         else:
-            # handle weather effects
-            current_weather = self.state.weather_schedule[self.state.real_env_steps]
-            current_weather = self.state.env_cfg.WEATHER_ID_TO_NAME[current_weather]
-            weather_cfg = weather.apply_weather(self.state, self.agents, current_weather)
-
             # 1. Check for malformed actions
-            
+
             if self.env_cfg.validate_action_space:
                 # This part is not absolutely necessary if you know for sure your actions are well formatted
                 for agent, unit_actions in actions.items():
-                    valid_acts, err_reason = self.action_space(agent).contains(unit_actions)
+                    valid_acts, err_reason = self.action_space(agent).contains(
+                        unit_actions
+                    )
                     if not valid_acts:
                         failed_agents[agent] = True
-                        self._log(f"{self.state.teams[agent]} Inappropriate action given. {err_reason}")
+                        self._log(
+                            f"{self.state.teams[agent]} Inappropriate action given. {err_reason}"
+                        )
 
             # we should except that actions is always of type dict, if not then erroring here is fine
             for agent, unit_actions in actions.items():
                 try:
                     for unit_id, action in unit_actions.items():
                         if "factory" in unit_id:
-                            self.state.factories[agent][unit_id].action_queue.append(format_factory_action(action))
+                            self.state.factories[agent][unit_id].action_queue.append(
+                                format_factory_action(action)
+                            )
                         elif "unit" in unit_id:
                             unit = self.state.units[agent][unit_id]
                             # if unit does not have more than ACTION_QUEUE_POWER_COST power, we skip updating the action queue and print warning
-                            update_power_req = self.state.env_cfg.ROBOTS[unit.unit_type.name].ACTION_QUEUE_POWER_COST * weather_cfg["power_loss_factor"]
+                            update_power_req = self.state.env_cfg.ROBOTS[
+                                unit.unit_type.name
+                            ].ACTION_QUEUE_POWER_COST
                             if unit.power < update_power_req:
-                                self._log(f"{agent} Tried to update action queue for {unit} requiring ({self.state.env_cfg.ROBOTS[unit.unit_type.name].ACTION_QUEUE_POWER_COST} x {weather_cfg['power_loss_factor']}) = {update_power_req} power but only had {unit.power} power. Power cost factor is {weather_cfg['power_loss_factor']} ")
+                                self._log(
+                                    f"{agent} Tried to update action queue for {unit} requiring {update_power_req}"
+                                )
                                 continue
                             formatted_actions = []
-                            if type(action) == list or (type(action) == np.ndarray and len(action.shape) == 2):
-                                trunked_actions = action[: self.env_cfg.UNIT_ACTION_QUEUE_SIZE]
-                                formatted_actions = [format_action_vec(a) for a in trunked_actions]
+                            if type(action) == list or (
+                                type(action) == np.ndarray and len(action.shape) == 2
+                            ):
+                                trunked_actions = action[
+                                    : self.env_cfg.UNIT_ACTION_QUEUE_SIZE
+                                ]
+                                formatted_actions = [
+                                    format_action_vec(a) for a in trunked_actions
+                                ]
                             else:
-                                self._log(f"{agent} Tried to update action queue for {unit} but did not provide an action queue, provided {action}")
+                                self._log(
+                                    f"{agent} Tried to update action queue for {unit} but did not provide an action queue, provided {action}"
+                                )
                                 failed_agents[agent] = True
                                 continue
                             unit.power -= update_power_req
-                            self.state.units[agent][unit_id].action_queue = formatted_actions
+                            self.state.units[agent][
+                                unit_id
+                            ].action_queue = formatted_actions
                 except Exception as e:
                     # catch errors when trying to format unit or factory actions
                     print(traceback.format_exc())
                     failed_agents[agent] = True
-        
+
             # 2. store actions by type
             actions_by_type: ActionsByType = defaultdict(list)
             for agent in self.agents:
-                if failed_agents[agent]: # skip failed agents
+                if failed_agents[agent]:  # skip failed agents
                     continue
                 for unit in self.state.units[agent].values():
                     unit_a: Action = unit.next_action()
@@ -624,37 +794,45 @@ def step(self, actions):
                         actions_by_type[unit_a.act_type].append((factory, unit_a))
 
             # 3. validate all actions against current state, throw away impossible actions
-            actions_by_type = validate_actions(self.env_cfg, self.state, actions_by_type, verbose=self.env_cfg.verbose, weather_cfg=weather_cfg)
+            actions_by_type = validate_actions(
+                self.env_cfg, self.state, actions_by_type, verbose=self.env_cfg.verbose
+            )
 
             if self.collect_stats:
                 lichen_before = self.state.board.lichen.copy()
                 lichen_strains_before = self.state.board.lichen_strains.copy()
 
-            self._handle_dig_actions(actions_by_type, weather_cfg)
+            self._handle_dig_actions(actions_by_type)
             self._handle_self_destruct_actions(actions_by_type)
-            self._handle_factory_build_actions(actions_by_type, weather_cfg)
-            self._handle_movement_actions(actions_by_type, weather_cfg)
+            self._handle_factory_build_actions(actions_by_type)
+            self._handle_movement_actions(actions_by_type)
             self._handle_recharge_actions(actions_by_type)
-            
+
             for agent in self.agents:
                 for factory in self.state.factories[agent].values():
                     # update information for lichen growing and cache it
                     factory.cache_water_info(self.state.board, self.env_cfg)
-            
+
             self._handle_factory_water_actions(actions_by_type)
             self._handle_transfer_actions(actions_by_type)
             self._handle_pickup_actions(actions_by_type)
-            
+
             # Update lichen
             self.state.board.lichen -= 1
-            self.state.board.lichen = self.state.board.lichen.clip(0, self.env_cfg.MAX_LICHEN_PER_TILE)
+            self.state.board.lichen = self.state.board.lichen.clip(
+                0, self.env_cfg.MAX_LICHEN_PER_TILE
+            )
             self.state.board.lichen_strains[self.state.board.lichen == 0] = -1
             if self.collect_stats:
                 lichen_change = self.state.board.lichen - lichen_before
                 for agent in self.agents:
                     for strain in self.state.teams[agent].factory_strains:
-                        start_of_step_lichen_tiles: np.ndarray = lichen_change[lichen_strains_before == strain]
-                        lichen_lost = start_of_step_lichen_tiles[start_of_step_lichen_tiles < 0].sum()
+                        start_of_step_lichen_tiles: np.ndarray = lichen_change[
+                            lichen_strains_before == strain
+                        ]
+                        lichen_lost = start_of_step_lichen_tiles[
+                            start_of_step_lichen_tiles < 0
+                        ].sum()
                         lichen_gained = start_of_step_lichen_tiles.sum() - lichen_lost
                         self.state.stats[agent]["lichen_grown"] += lichen_gained
                         self.state.stats[agent]["lichen_lost"] -= lichen_lost
@@ -676,13 +854,18 @@ def step(self, actions):
             if is_day(self.env_cfg, self.state.real_env_steps):
                 for agent in self.agents:
                     for u in self.state.units[agent].values():
-                        u.power = u.power + math.ceil(self.env_cfg.ROBOTS[u.unit_type.name].CHARGE * weather_cfg["power_gain_factor"])
+                        u.power = u.power + self.env_cfg.ROBOTS[u.unit_type.name].CHARGE
                         u.power = min(u.power, u.unit_cfg.BATTERY_CAPACITY)
             for agent in self.agents:
                 for f in self.state.factories[agent].values():
-                    # Factories are immune to weather thanks to using nuclear reactors instead
-                    f.power = f.power + math.ceil(self.env_cfg.FACTORY_CHARGE)
-
+                    # natural nuclear energy generation
+                    f.power = f.power + self.env_cfg.FACTORY_CHARGE
+                    # lichen/plant power
+                    f.power = (
+                        f.power
+                        + len(f.grow_lichen_positions)
+                        * self.env_cfg.POWER_PER_CONNECTED_LICHEN_TILE
+                    )
 
         # always set rubble under factories to 0.
         self.state.board.rubble[self.state.board.factory_occupancy_map != -1] = 0
@@ -699,7 +882,9 @@ def step(self, actions):
                 if failed_agents[agent]:
                     rewards[agent] = -1000
                 else:
-                    agent_lichen_mask = np.isin(self.state.board.lichen_strains, strain_ids)
+                    agent_lichen_mask = np.isin(
+                        self.state.board.lichen_strains, strain_ids
+                    )
                     rewards[agent] = self.state.board.lichen[agent_lichen_mask].sum()
             else:
                 # if this was not initialize then agent failed in step 0
@@ -709,7 +894,9 @@ def step(self, actions):
         self.env_steps += 1
         self.state.env_steps += 1
         env_done = self.state.real_env_steps >= self.state.env_cfg.max_episode_length
-        env_done = env_done or failed_agents["player_0"] or failed_agents["player_1"] # env is done if any agent fails.
+        env_done = (
+            env_done or failed_agents["player_0"] or failed_agents["player_1"]
+        )  # env is done if any agent fails.
         dones = {agent: env_done or failed_agents[agent] for agent in self.agents}
 
         # generate observations
@@ -728,7 +915,12 @@ def step(self, actions):
 
     ### Game Logic ###
     def add_unit(self, team: Team, unit_type, pos: np.ndarray):
-        unit = Unit(team=team, unit_type=unit_type, unit_id=f"unit_{self.state.global_id}", env_cfg=self.env_cfg)
+        unit = Unit(
+            team=team,
+            unit_type=unit_type,
+            unit_id=f"unit_{self.state.global_id}",
+            env_cfg=self.env_cfg,
+        )
         unit.pos.pos = pos.copy()
         self.state.global_id += 1
         self.state.units[team.agent][unit.unit_id] = unit
@@ -736,21 +928,29 @@ def add_unit(self, team: Team, unit_type, pos: np.ndarray):
         return unit
 
     def add_factory(self, team: Team, pos: np.ndarray):
-        factory = Factory(team=team, unit_id=f"factory_{self.state.global_id}", num_id=self.state.global_id)
+        factory = Factory(
+            team=team,
+            unit_id=f"factory_{self.state.global_id}",
+            num_id=self.state.global_id,
+        )
         factory.pos.pos = list(pos)
         factory.cargo.water = self.env_cfg.INIT_WATER_METAL_PER_FACTORY
         factory.cargo.metal = self.env_cfg.INIT_WATER_METAL_PER_FACTORY
         factory.power = self.env_cfg.INIT_POWER_PER_FACTORY
         if self.state.board.valid_spawns_mask[pos[0], pos[1]] == 0:
             # Check if any tiles under the factory are invalid spawn tile.
-            # TODO - min distance between Factories? stone: I think it's a bad strategy to try and enclose a opponent factory anyway, 
+            # TODO - min distance between Factories? stone: I think it's a bad strategy to try and enclose a opponent factory anyway,
             # wastes a few factories and hard to maintain
-            self._log(f"{team.agent} cannot place factory at {pos[0]}, {pos[1]} as it overlaps an existing factory or is on top of a resource")
+            self._log(
+                f"{team.agent} cannot place factory at {pos[0]}, {pos[1]} as it overlaps an existing factory or is on top of a resource"
+            )
             return None
         if self.state.board.factory_occupancy_map[factory.pos_slice].max() >= 0:
-            self._log(f"{team.agent} cannot overlap factory placement. Existing factory at {factory.pos} already.")
+            self._log(
+                f"{team.agent} cannot overlap factory placement. Existing factory at {factory.pos} already."
+            )
             return None
-        
+
         self.state.teams[team.agent].factory_strains += [factory.num_id]
 
         self.state.factories[team.agent][factory.unit_id] = factory
@@ -759,10 +959,15 @@ def add_factory(self, team: Team, pos: np.ndarray):
         invalid_spawn_indices = factory.min_dist_slice
         # TODO: perf - this can be faster
         # self.state.board.valid_spawns_mask[factory.pos_slice] = False
-        for x,y in invalid_spawn_indices:
-            if x < 0 or y < 0 or x >= self.state.board.rubble.shape[0] or y >= self.state.board.rubble.shape[1]:
+        for x, y in invalid_spawn_indices:
+            if (
+                x < 0
+                or y < 0
+                or x >= self.state.board.rubble.shape[0]
+                or y >= self.state.board.rubble.shape[1]
+            ):
                 continue
-            self.state.board.valid_spawns_mask[x,y] = False
+            self.state.board.valid_spawns_mask[x, y] = False
         self.state.board.rubble[factory.pos_slice] = 0
         self.state.board.ice[factory.pos_slice] = 0
         self.state.board.ore[factory.pos_slice] = 0
@@ -775,7 +980,8 @@ def destroy_unit(self, unit: Unit):
         # NOTE this doesn't remove unit reference from board map
         """
         self.state.board.rubble[unit.pos.x, unit.pos.y] = min(
-            self.state.board.rubble[unit.pos.x, unit.pos.y] + unit.unit_cfg.RUBBLE_AFTER_DESTRUCTION,
+            self.state.board.rubble[unit.pos.x, unit.pos.y]
+            + unit.unit_cfg.RUBBLE_AFTER_DESTRUCTION,
             self.env_cfg.MAX_RUBBLE,
         )
         self.state.board.lichen[unit.pos.x, unit.pos.y] = 0
@@ -784,12 +990,17 @@ def destroy_unit(self, unit: Unit):
 
     def destroy_factory(self, factory: Factory):
         # spray rubble on every factory tile
-        self.state.board.rubble[factory.pos_slice] += self.env_cfg.FACTORY_RUBBLE_AFTER_DESTRUCTION
-        self.state.board.rubble[factory.pos_slice] = self.state.board.rubble[factory.pos_slice].clip(0, self.env_cfg.MAX_RUBBLE)
+        self.state.board.rubble[
+            factory.pos_slice
+        ] += self.env_cfg.FACTORY_RUBBLE_AFTER_DESTRUCTION
+        self.state.board.rubble[factory.pos_slice] = self.state.board.rubble[
+            factory.pos_slice
+        ].clip(0, self.env_cfg.MAX_RUBBLE)
         self.state.board.factory_occupancy_map[factory.pos_slice] = -1
         del self.state.factories[factory.team.agent][factory.unit_id]
         del self.state.board.factory_map[self.state.board.pos_hash(factory.pos)]
 
+
 def raw_env() -> LuxAI_S2:
     """
     To support the AEC API, the raw_env() function just uses the from_parallel
diff --git a/luxai_s2/luxai_s2/factory.py b/luxai_s2/luxai_s2/factory.py
index ea80bf28..a5bfe302 100644
--- a/luxai_s2/luxai_s2/factory.py
+++ b/luxai_s2/luxai_s2/factory.py
@@ -1,23 +1,35 @@
 from __future__ import annotations
+
+import time
 from collections import deque
 from itertools import product
-import time
 from typing import List
+
 import numpy as np
+
+from luxai_s2.actions import move_deltas
 from luxai_s2.config import EnvConfig
+from luxai_s2.globals import TERM_COLORS
 from luxai_s2.map.board import Board
 from luxai_s2.map.position import Position
 from luxai_s2.team import Team
 from luxai_s2.unit import UnitCargo
-from luxai_s2.actions import move_deltas
-from luxai_s2.globals import TERM_COLORS
+
 try:
     from termcolor import colored
 except:
     pass
 
 
-def compute_water_info(init: np.ndarray, MIN_LICHEN_TO_SPREAD: int, lichen: np.ndarray, lichen_strains: np.ndarray, factory_occupancy_map: np.ndarray, strain_id: int, forbidden: np.ndarray):
+def compute_water_info(
+    init: np.ndarray,
+    MIN_LICHEN_TO_SPREAD: int,
+    lichen: np.ndarray,
+    lichen_strains: np.ndarray,
+    factory_occupancy_map: np.ndarray,
+    strain_id: int,
+    forbidden: np.ndarray,
+):
     # TODO - improve the performance here with cached solution
     frontier = deque(init)
     seen = set(map(tuple, init))
@@ -30,7 +42,12 @@ def compute_water_info(init: np.ndarray, MIN_LICHEN_TO_SPREAD: int, lichen: np.n
             print("Error! Lichen Growth calculation took too long")
             break
         pos = frontier.popleft()
-        if pos[0] < 0 or pos[1] < 0 or pos[0] >= forbidden.shape[0] or pos[1] >= forbidden.shape[1]:
+        if (
+            pos[0] < 0
+            or pos[1] < 0
+            or pos[0] >= forbidden.shape[0]
+            or pos[1] >= forbidden.shape[1]
+        ):
             continue
 
         if forbidden[pos[0], pos[1]]:
@@ -42,14 +59,21 @@ def compute_water_info(init: np.ndarray, MIN_LICHEN_TO_SPREAD: int, lichen: np.n
         for move_delta in move_deltas[1:]:
             check_pos = pos + move_delta
             # check surrounding tiles on the map
-            if check_pos[0] < 0 or check_pos[1] < 0 or check_pos[0] >= H or check_pos[1] >= W: continue
+            if (
+                check_pos[0] < 0
+                or check_pos[1] < 0
+                or check_pos[0] >= H
+                or check_pos[1] >= W
+            ):
+                continue
 
-            # If any neighbor 1. has a different strain, or 2. is a different factory, 
+            # If any neighbor 1. has a different strain, or 2. is a different factory,
             # then the current pos cannot grow
             adj_strain = lichen_strains[check_pos[0], check_pos[1]]
             adj_factory = factory_occupancy_map[check_pos[0], check_pos[1]]
-            if (adj_strain != -1 and adj_strain != strain_id) \
-                or (adj_factory != -1 and adj_factory != strain_id):
+            if (adj_strain != -1 and adj_strain != strain_id) or (
+                adj_factory != -1 and adj_factory != strain_id
+            ):
                 can_grow = False
 
             # if seen, skip
@@ -59,8 +83,9 @@ def compute_water_info(init: np.ndarray, MIN_LICHEN_TO_SPREAD: int, lichen: np.n
             # we add it to the frontier only in two cases:
             #  1. it is an empty tile, and current pos has enough lichen to expand.
             #  2. both current tile and check_pos are of our strain.
-            if (adj_strain == -1 and pos_lichen >= MIN_LICHEN_TO_SPREAD) \
-                or (adj_strain == strain_id and pos_strain == strain_id):
+            if (adj_strain == -1 and pos_lichen >= MIN_LICHEN_TO_SPREAD) or (
+                adj_strain == strain_id and pos_strain == strain_id
+            ):
                 seen.add(tuple(check_pos))
                 frontier.append(check_pos)
 
@@ -68,6 +93,7 @@ def compute_water_info(init: np.ndarray, MIN_LICHEN_TO_SPREAD: int, lichen: np.n
             grow_lichen_positions.add((pos[0], pos[1]))
     return grow_lichen_positions
 
+
 class Factory:
     def __init__(self, team: Team, unit_id: str, num_id: int) -> None:
         self.team_id = team.team_id
@@ -78,18 +104,107 @@ def __init__(self, team: Team, unit_id: str, num_id: int) -> None:
         self.cargo = UnitCargo()
         self.num_id = num_id
         self.action_queue = []
-        self.grow_lichen_positions = None
+        self.grow_lichen_positions = set()
 
     @property
     def pos_slice(self):
-        return slice(self.pos.x - 1, self.pos.x + 2), slice(self.pos.y - 1, self.pos.y + 2)
+        return slice(self.pos.x - 1, self.pos.x + 2), slice(
+            self.pos.y - 1, self.pos.y + 2
+        )
 
     @property
     def min_dist_slice(self):
-        deltas = np.array([(4, 0), (5, 1), (0, -4), (0, 5), (2, 2), (-4, 1), (-1, -1), (-1, -2), (-2, -1), (-2, -2), (4, 2), (0, -2), (0, -1), (2, 4), (-1, 0), (-2, 0), (-3, 1), (0, 0), (2, -3), (-2, 2), (-1, 2), (3, 1), (5, -1), (-3, 3), (0, 2), (1, 3), (-4, -1), (-1, -5), (-4, -2), (-1, 4), (-2, 4), (3, 3), (5, 0), (-6, 0), (0, -5), (1, -4), (1, 5), (-4, 0), (-1, -3), (-2, -3), (-5, 1), (-3, -1), (0, -3), (-3, -2), (1, -1), (1, -2), (-4, 2), (3, -1), (3, -2), (-3, 0), (1, 0), (3, 0), (-3, 2), (1, 2), (0, 4), (2, 1), (3, 2), (4, 1), (-5, -1), (1, -5), (1, 4), (0, 6), (2, 3), (-1, -4), (-5, 0), (-3, -3), (1, -3), (2, -4), (-1, 1), (-2, 1), (3, -3), (0, 1), (2, -1), (2, -2), (-1, 3), (-2, 3), (4, -1), (4, -2), (0, -6), (1, 1), (0, 3), (2, 0), (6, 0), (-2, -4), (-1, 5)])
+        deltas = np.array(
+            [
+                (4, 0),
+                (5, 1),
+                (0, -4),
+                (0, 5),
+                (2, 2),
+                (-4, 1),
+                (-1, -1),
+                (-1, -2),
+                (-2, -1),
+                (-2, -2),
+                (4, 2),
+                (0, -2),
+                (0, -1),
+                (2, 4),
+                (-1, 0),
+                (-2, 0),
+                (-3, 1),
+                (0, 0),
+                (2, -3),
+                (-2, 2),
+                (-1, 2),
+                (3, 1),
+                (5, -1),
+                (-3, 3),
+                (0, 2),
+                (1, 3),
+                (-4, -1),
+                (-1, -5),
+                (-4, -2),
+                (-1, 4),
+                (-2, 4),
+                (3, 3),
+                (5, 0),
+                (-6, 0),
+                (0, -5),
+                (1, -4),
+                (1, 5),
+                (-4, 0),
+                (-1, -3),
+                (-2, -3),
+                (-5, 1),
+                (-3, -1),
+                (0, -3),
+                (-3, -2),
+                (1, -1),
+                (1, -2),
+                (-4, 2),
+                (3, -1),
+                (3, -2),
+                (-3, 0),
+                (1, 0),
+                (3, 0),
+                (-3, 2),
+                (1, 2),
+                (0, 4),
+                (2, 1),
+                (3, 2),
+                (4, 1),
+                (-5, -1),
+                (1, -5),
+                (1, 4),
+                (0, 6),
+                (2, 3),
+                (-1, -4),
+                (-5, 0),
+                (-3, -3),
+                (1, -3),
+                (2, -4),
+                (-1, 1),
+                (-2, 1),
+                (3, -3),
+                (0, 1),
+                (2, -1),
+                (2, -2),
+                (-1, 3),
+                (-2, 3),
+                (4, -1),
+                (4, -2),
+                (0, -6),
+                (1, 1),
+                (0, 3),
+                (2, 0),
+                (6, 0),
+                (-2, -4),
+                (-1, 5),
+            ]
+        )
         return self.pos.pos + deltas
 
-
     def refine_step(self, config: EnvConfig):
         max_consumed_ice = min(self.cargo.ice, config.FACTORY_PROCESSING_RATE_WATER)
         max_consumed_ore = min(self.cargo.ore, config.FACTORY_PROCESSING_RATE_METAL)
@@ -117,17 +232,48 @@ def cache_water_info(self, board: Board, env_cfg: EnvConfig):
            x x x
 
         """
-        forbidden = (board.rubble > 0) | (board.factory_occupancy_map != -1) | (board.ice > 0) | (board.ore > 0)
-        deltas = [np.array([0, -2]),  np.array([-1, -2]),  np.array([1, -2]),  np.array([0, 2]),  np.array([-1, 2]),  np.array([1, 2]),
-                    np.array([2, 0]),  np.array([2, -1]),  np.array([2, 1]),  np.array([-2, 0]),  np.array([-2, -1]),  np.array([-2, 1])]
+        forbidden = (
+            (board.rubble > 0)
+            | (board.factory_occupancy_map != -1)
+            | (board.ice > 0)
+            | (board.ore > 0)
+        )
+        deltas = [
+            np.array([0, -2]),
+            np.array([-1, -2]),
+            np.array([1, -2]),
+            np.array([0, 2]),
+            np.array([-1, 2]),
+            np.array([1, 2]),
+            np.array([2, 0]),
+            np.array([2, -1]),
+            np.array([2, 1]),
+            np.array([-2, 0]),
+            np.array([-2, -1]),
+            np.array([-2, 1]),
+        ]
         init_arr = np.stack(deltas) + self.pos.pos
-        self.grow_lichen_positions = compute_water_info(init_arr, env_cfg.MIN_LICHEN_TO_SPREAD, board.lichen, board.lichen_strains, board.factory_occupancy_map, self.num_id, forbidden)
+        self.grow_lichen_positions = compute_water_info(
+            init_arr,
+            env_cfg.MIN_LICHEN_TO_SPREAD,
+            board.lichen,
+            board.lichen_strains,
+            board.factory_occupancy_map,
+            self.num_id,
+            forbidden,
+        )
+
     def water_cost(self, config: EnvConfig):
-        return int(np.ceil(len(self.grow_lichen_positions) / config.LICHEN_WATERING_COST_FACTOR))
+        return int(
+            np.ceil(
+                len(self.grow_lichen_positions) / config.LICHEN_WATERING_COST_FACTOR
+            )
+        )
 
     ### Add and sub resource functions copied over from unit.py code, can we consolidate them somewhere?
     def add_resource(self, resource_id, transfer_amount):
-        if transfer_amount < 0: transfer_amount = 0
+        if transfer_amount < 0:
+            transfer_amount = 0
         if resource_id == 0:
             self.cargo.ice += transfer_amount
         elif resource_id == 1:
@@ -139,9 +285,11 @@ def add_resource(self, resource_id, transfer_amount):
         elif resource_id == 4:
             self.power += transfer_amount
         return int(transfer_amount)
+
     def sub_resource(self, resource_id, amount):
         # subtract/transfer out as much as you min(have, request)
-        if amount < 0: amount = 0
+        if amount < 0:
+            amount = 0
         if resource_id == 0:
             transfer_amount = min(self.cargo.ice, amount)
             self.cargo.ice -= transfer_amount
@@ -159,18 +307,14 @@ def sub_resource(self, resource_id, amount):
             self.power -= transfer_amount
         return int(transfer_amount)
 
-
-
-
     def state_dict(self):
         return dict(
             pos=self.pos.pos,
             power=self.power,
             cargo=self.cargo.state_dict(),
             unit_id=self.unit_id,
-            strain_id=self.num_id, # number version of unit_id
-            team_id=self.team_id
-
+            strain_id=self.num_id,  # number version of unit_id
+            team_id=self.team_id,
         )
 
     def __str__(self) -> str:
diff --git a/luxai_s2/luxai_s2/globals.py b/luxai_s2/luxai_s2/globals.py
index f3645a98..5313c18b 100644
--- a/luxai_s2/luxai_s2/globals.py
+++ b/luxai_s2/luxai_s2/globals.py
@@ -1,9 +1,13 @@
 import os
+
 TERM_COLORS = True
 try:
     from termcolor import colored
-    TERM_COLORS = os.environ["LUX_COLORS"] == 'False' if "LUX_COLORS" in os.environ else True
+
+    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
\ No newline at end of file
+    pass
diff --git a/luxai_s2/luxai_s2/map/__init__.py b/luxai_s2/luxai_s2/map/__init__.py
index bee299a6..e69de29b 100644
--- a/luxai_s2/luxai_s2/map/__init__.py
+++ b/luxai_s2/luxai_s2/map/__init__.py
@@ -1,3 +0,0 @@
-from .board import Board
-from .position import Position
-from .weather import generate_weather_schedule, apply_weather
\ No newline at end of file
diff --git a/luxai_s2/luxai_s2/map/bfs_deltas_gen.py b/luxai_s2/luxai_s2/map/bfs_deltas_gen.py
index bd3d6841..24bd4809 100644
--- a/luxai_s2/luxai_s2/map/bfs_deltas_gen.py
+++ b/luxai_s2/luxai_s2/map/bfs_deltas_gen.py
@@ -1,5 +1,8 @@
 import numpy as np
+
 move_deltas = np.array([[0, 0], [0, -1], [1, 0], [0, 1], [-1, 0]])
+
+
 def gen_deltas(size=3):
     start = (0, 0)
     f = [start]
@@ -12,9 +15,10 @@ def gen_deltas(size=3):
             break
         r.add(pos)
         for md in move_deltas:
-            new_pos = (pos[0] + md[0],pos[1] + md[1])
+            new_pos = (pos[0] + md[0], pos[1] + md[1])
             if new_pos not in seen:
                 f.append(new_pos)
     print(r)
-gen_deltas(6)
 
+
+gen_deltas(6)
diff --git a/luxai_s2/luxai_s2/map/board.py b/luxai_s2/luxai_s2/map/board.py
index 99b3220e..369c93ec 100644
--- a/luxai_s2/luxai_s2/map/board.py
+++ b/luxai_s2/luxai_s2/map/board.py
@@ -1,12 +1,16 @@
 from __future__ import annotations
+
 from collections import defaultdict
-from typing import Dict, List
+from typing import TYPE_CHECKING, Dict, List
+
 import numpy as np
-from typing import TYPE_CHECKING
+
 from luxai_s2.config import EnvConfig
+
 if TYPE_CHECKING:
     from luxai_s2.factory import Factory
     from luxai_s2.state import State
+
 from luxai_s2.map.position import Position
 from luxai_s2.map_generator.generator import GameMap
 from luxai_s2.unit import Unit
@@ -15,14 +19,31 @@
 class Board:
     def __init__(self, seed=None, env_cfg: EnvConfig = None) -> None:
         self.env_cfg = env_cfg
-        self.height =  self.env_cfg.map_size
-        self.width =  self.env_cfg.map_size
+        self.height = self.env_cfg.map_size
+        self.width = self.env_cfg.map_size
         self.seed = seed
         rng = np.random.RandomState(seed=seed)
         map_type = rng.choice(["Cave", "Mountain"])
+        map_distribution_type = rng.choice(
+            [
+                "high_ice_high_ore",
+                "high_ice_low_ore",
+                "low_ice_high_ore",
+                "low_ice_low_ore",
+            ]
+        )
         symmetry = rng.choice(["horizontal", "vertical"])
-        self.map = GameMap.random_map(seed=seed, symmetry=symmetry, map_type=map_type, width=self.width, height=self.height)
-        self.factories_per_team = rng.randint(env_cfg.MIN_FACTORIES, env_cfg.MAX_FACTORIES + 1)
+        self.map = GameMap.random_map(
+            seed=seed,
+            symmetry=symmetry,
+            map_type=map_type,
+            map_distribution_type=map_distribution_type,
+            width=self.width,
+            height=self.height,
+        )
+        self.factories_per_team = rng.randint(
+            env_cfg.MIN_FACTORIES, env_cfg.MAX_FACTORIES + 1
+        )
 
         # remove bottom once generator is fully ready
         self.map.rubble = self.map.rubble.astype(int)
@@ -38,39 +59,54 @@ def __init__(self, seed=None, env_cfg: EnvConfig = None) -> None:
         # TODO consider unit occupancy map, may speed up computation by storing more info in a form indexable in numpy arrays
 
         # maps center of factory to the factory
-        self.factory_map: Dict[str, 'Factory'] = dict()
+        self.factory_map: Dict[str, "Factory"] = dict()
         # != -1 if a factory tile is on the location. Equals factory number id
         self.factory_occupancy_map = -np.ones((self.height, self.width), dtype=int)
 
-
         # Valid spawn locations
         self.valid_spawns_mask = np.ones((self.height, self.width), dtype=bool)
         resource_mask = (self.map.ice != 0) | (self.map.ore != 0)
         # generate all 9 shifts of the resource mask. Below is a simple solution to add 1s around any existing 1s in resource_mask
         init_mask = np.zeros((self.height + 2, self.width + 2), dtype=bool)
-        for delta in np.array([[0,0],[0, 1], [0, -1], [1, 1], [1,0], [1, -1], [-1, -1], [-1, 0], [-1, 1]]):
+        for delta in np.array(
+            [
+                [0, 0],
+                [0, 1],
+                [0, -1],
+                [1, 1],
+                [1, 0],
+                [1, -1],
+                [-1, -1],
+                [-1, 0],
+                [-1, 1],
+            ]
+        ):
             s0 = delta[0] + 1
             s1 = delta[1] + 1
             e0 = -1 + delta[0]
             e1 = -1 + delta[1]
-            if e0 == 0: e0 = None
-            if e1 == 0: e1 = None
+            if e0 == 0:
+                e0 = None
+            if e1 == 0:
+                e1 = None
             init_mask[s0:e0, s1:e1] = init_mask[s0:e0, s1:e1] | resource_mask
-        resource_mask = init_mask[1:-1,1:-1]
+        resource_mask = init_mask[1:-1, 1:-1]
         self.valid_spawns_mask[resource_mask] = False
         self.valid_spawns_mask[0] = False
         self.valid_spawns_mask[-1] = False
-        self.valid_spawns_mask[:,0] = False
-        self.valid_spawns_mask[:,-1] = False
+        self.valid_spawns_mask[:, 0] = False
+        self.valid_spawns_mask[:, -1] = False
         # grow the 0s by 1 radius
 
     def pos_hash(self, pos: Position):
         return f"{pos.x},{pos.y}"
+
     def get_units_at(self, pos: Position):
         pos_hash = self.pos_hash(pos)
         if pos_hash in self.units_map:
             return self.units_map[pos_hash]
         return None
+
     def get_factory_at(self, state: State, pos: Position):
         f_id = self.factory_occupancy_map[pos.x, pos.y]
         if f_id != -1:
@@ -84,12 +120,15 @@ def get_factory_at(self, state: State, pos: Position):
     @property
     def rubble(self) -> np.ndarray:
         return self.map.rubble
+
     @property
     def ice(self) -> np.ndarray:
         return self.map.ice
+
     @property
     def ore(self):
         return self.map.ore
+
     def state_dict(self):
         return dict(
             rubble=self.rubble.copy(),
diff --git a/luxai_s2/luxai_s2/map/weather.py b/luxai_s2/luxai_s2/map/weather.py
deleted file mode 100644
index 78ff0f48..00000000
--- a/luxai_s2/luxai_s2/map/weather.py
+++ /dev/null
@@ -1,47 +0,0 @@
-from typing import List
-import numpy as np
-from luxai_s2.config import EnvConfig
-from luxai_s2.state.state import State
-from luxai_s2.unit import Unit
-def generate_weather_schedule(rng: np.random.RandomState, cfg: EnvConfig):
-    # randomly generate 3-5 events, each lasting 20 turns
-    # no event can overlap another
-    num_events = rng.randint(cfg.NUM_WEATHER_EVENTS_RANGE[0], cfg.NUM_WEATHER_EVENTS_RANGE[1] + 1)
-    last_event_end_step = 0
-    # TODO - make a smarter algorithm to speed up the generation here
-    available_times = set(list(range(cfg.max_episode_length - 30)))
-    schedule = np.zeros(cfg.max_episode_length, dtype=int)
-    for i in range(num_events):
-        weather_id = rng.randint(1, len(cfg.WEATHER_ID_TO_NAME))
-        weather = cfg.WEATHER_ID_TO_NAME[weather_id]
-        weather_cfg = cfg.WEATHER[weather]
-        weather_time_range = weather_cfg["TIME_RANGE"]
-        # use rejection sampling
-        while True:
-            start_time = rng.randint(0, cfg.max_episode_length - 20)
-            duration = rng.randint(weather_time_range[0], weather_time_range[1] + 1)
-            requested_times = set(list(range(start_time, start_time + duration)))
-            if requested_times.issubset(available_times):
-                available_times.difference_update(requested_times)
-                schedule[start_time: start_time + duration] = weather_id
-                break
-    return schedule
-    
-
-def apply_weather(state: State, agents: List[str], current_weather):
-    if current_weather == "MARS_QUAKE":
-        for agent in agents:
-            for unit in state.units[agent].values():
-                unit: Unit
-                old_rubble = state.board.rubble[unit.pos.x, unit.pos.y]
-                if state.board.factory_occupancy_map[unit.pos.x, unit.pos.y] == -1:
-                    state.board.rubble[unit.pos.x, unit.pos.y] = min(state.env_cfg.MAX_RUBBLE, old_rubble + state.env_cfg.ROBOTS[unit.unit_type.name].RUBBLE_AFTER_DESTRUCTION)
-                    state.board.lichen[unit.pos.x, unit.pos.y] = 0
-        return dict(power_gain_factor=1, power_loss_factor=1)
-    elif current_weather == "COLD_SNAP":
-        return dict(power_gain_factor=1, power_loss_factor=state.env_cfg.WEATHER["COLD_SNAP"]["POWER_CONSUMPTION"])
-    elif current_weather == "DUST_STORM":
-        return dict(power_gain_factor=state.env_cfg.WEATHER["DUST_STORM"]["POWER_GAIN"], power_loss_factor=1)
-    elif current_weather == "SOLAR_FLARE":
-        return dict(power_gain_factor=state.env_cfg.WEATHER["SOLAR_FLARE"]["POWER_GAIN"], power_loss_factor=1)
-    return dict(power_gain_factor=1, power_loss_factor=1)
diff --git a/luxai_s2/luxai_s2/map_generator/generator.py b/luxai_s2/luxai_s2/map_generator/generator.py
index ed4a92f7..df48fcfa 100644
--- a/luxai_s2/luxai_s2/map_generator/generator.py
+++ b/luxai_s2/luxai_s2/map_generator/generator.py
@@ -1,8 +1,13 @@
+from dataclasses import dataclass
+from typing import Tuple
+
 import numpy as np
-from scipy.ndimage import convolve, maximum_filter
 from scipy.fft import dctn, idctn
-from luxai_s2.map_generator.visualize import viz
+from scipy.ndimage import convolve, maximum_filter
+
 from luxai_s2.map_generator.symnoise import SymmetricNoise, symmetrize
+from luxai_s2.map_generator.visualize import viz
+
 
 class GameMap(object):
     def __init__(self, rubble, ice, ore, symmetry):
@@ -28,7 +33,14 @@ def noise(seed=None, noise=None, noise_shift=None, symmetry=None):
         return noise
 
     @staticmethod
-    def random_map(seed=None, map_type=None, symmetry=None, width=None, height=None):
+    def random_map(
+        seed=None,
+        map_type=None,
+        map_distribution_type=None,
+        symmetry=None,
+        width=None,
+        height=None,
+    ):
         """random_map
         Returns a random map. Can specify seed.
         """
@@ -41,26 +53,75 @@ def random_map(seed=None, map_type=None, symmetry=None, width=None, height=None)
         if height is None:
             height = width
         if not map_type:
-            map_type = noise.random.choice(["Cave", "Craters", "Island", "Mountain"])
+            map_type = noise.random.choice(["Cave", "Mountain"])
+        if map_type == "Cave":
+            config_registry = CAVE_CONFIGS
+        elif map_type == "Mountain":
+            config_registry = MOUNTAIN_CONFIGS
+        if not map_distribution_type:
+            map_distribution_type = noise.random.choice(
+                [
+                    "high_ice_high_ore",
+                    "high_ice_low_ore",
+                    "low_ice_high_ore",
+                    "low_ice_low_ore",
+                ]
+            )
         if not symmetry:
-            symmetry = noise.random.choice(["horizontal", "vertical", "rotational", "/", "\\"])
+            symmetry = noise.random.choice(
+                ["horizontal", "vertical", "rotational", "/", "\\"]
+            )
         noise.update_symmetry(symmetry)
-        return eval(map_type)(width, height, symmetry, noise=noise)
+        return eval(map_type)(
+            width,
+            height,
+            symmetry,
+            noise=noise,
+            config=config_registry[map_distribution_type],
+        )
+
+
+@dataclass
+class CaveConfig:
+    # default parameters -> {'rubble_count': 48.165737847222225, 'ice': 37.93, 'ore': 38.36}
+    ice_high_range: Tuple[float, float] = (99, 100)
+    ice_mid_range: Tuple[float, float] = (91, 91.7)
+
+    ore_high_range: Tuple[float, float] = (98.7, 100)
+
+    # deeper in walls
+    ore_mid_range: Tuple[float, float] = (81, 81.5)
+
+
+CAVE_CONFIGS = dict(
+    # {'rubble_count': 48.165737847222225, 'ice': 37.93, 'ore': 38.36}
+    high_ice_high_ore=CaveConfig(),
+    low_ice_high_ore=CaveConfig(ice_high_range=(99.7, 100), ice_mid_range=(91, 91.5)),
+    high_ice_low_ore=CaveConfig(
+        ore_high_range=(99.6, 100),
+        ore_mid_range=(81.1, 81.4),
+    ),
+    # {'rubble_count': 48.165737847222225, 'ice': 16.96, 'ore': 14.57}
+    low_ice_low_ore=CaveConfig(
+        ore_high_range=(99.6, 100),
+        ore_mid_range=(81.1, 81.4),
+        ice_high_range=(99.7, 100),
+        ice_mid_range=(91, 91.5),
+    ),
+)
 
-    @staticmethod
-    def random_maps(n, seed=None):
-        noise = GameMap.noise(seed)
-        for i in range(n):
-            noise = GameMap.noise(noise, noise_shift=100)
-            width = height = noise.random.randint(32, 64)
-            map_type = noise.random.choice(["Cave", "Craters", "Island", "Mountain"])
-            symmetry = noise.random.choice(["horizontal", "vertical", "rotational", "/", "\\"])
-            noise.update_symmetry(symmetry)
-
-            yield eval(map_type)(width, height, symmetry, noise=noise)
 
 class Cave(GameMap):
-    def __init__(self, width=64, height=64, symmetry="vertical", seed=None, noise=None, noise_shift=None):
+    def __init__(
+        self,
+        width=64,
+        height=64,
+        symmetry="vertical",
+        seed=None,
+        noise=None,
+        noise_shift=None,
+        config: CaveConfig = CAVE_CONFIGS["high_ice_high_ore"],
+    ):
         """Cave
         Builds a cave system of size width x height and symmetrical across the `symmetry` axis.
 
@@ -83,10 +144,10 @@ def __init__(self, width=64, height=64, symmetry="vertical", seed=None, noise=No
         mask = noise.random.randint(3, size=(height, width))
         mask[mask >= 1] = 1
         # symmetrize(mask, symmetry)
-        
+
         # Build clumps of ones (will be interior of caves)
         for i in range(3):
-            mask = convolve(mask, [[1]*3]*3, mode="constant", cval=0) // 6
+            mask = convolve(mask, [[1] * 3] * 3, mode="constant", cval=0) // 6
         mask = 1 - mask
 
         # Create cave wall
@@ -99,30 +160,45 @@ def __init__(self, width=64, height=64, symmetry="vertical", seed=None, noise=No
         y = np.linspace(0, 1, height)
         rubble = noise(x, y) * 50 + 50
         rubble = rubble.round()
-        rubble[mask==1] //= 5 # Cave walls
-        rubble[mask==0] = 0 # Interior of cave
+        rubble[mask == 1] //= 5  # Cave walls
+        rubble[mask == 0] = 0  # Interior of cave
 
         # Make some noisy ice, most ice is on cave edges
         ice = noise(x, y + 100)
         ice[mask > 1] = 0
         ice[mask == 0] = 0
-        mid_mask = (ice < np.percentile(ice, 92)) & (ice > np.percentile(ice, 91))
-        ice[ice < np.percentile(ice, 98)] = 0
-        ice[ice != 0] = 1
-        ice[mid_mask] = 1
+        mid_mask = (ice > np.percentile(ice, config.ice_mid_range[0])) & (
+            ice < np.percentile(ice, config.ice_mid_range[1])
+        )
+        high_mask = (ice > np.percentile(ice, config.ice_high_range[0])) & (
+            ice < np.percentile(ice, config.ice_high_range[1])
+        )
+        ice = mid_mask | high_mask
 
         # Make some noisy ore, most ore is outside caves
         ore = noise(x, y - 100)
-        ore[mask==1] = 0
-        ore[mask==0] = 0
-        mid_mask = (ore < np.percentile(ore, 92)) & (ore > np.percentile(ore, 91))
-        ore[(ore < np.percentile(ore, 98))] = 0
-        ore[ore != 0] = 1
-        ore[mid_mask] = 1
+        ore[mask == 1] = 0
+        ore[mask == 0] = 0
+        mid_mask = (ore > np.percentile(ore, config.ore_mid_range[0])) & (
+            ore < np.percentile(ore, config.ore_mid_range[1])
+        )
+        high_mask = (ore > np.percentile(ore, config.ore_high_range[0])) & (
+            ore < np.percentile(ore, config.ore_high_range[1])
+        )
+        ore = mid_mask | high_mask
         super().__init__(rubble, ice, ore, symmetry)
 
+
 class Craters(GameMap):
-    def __init__(self, width=64, height=64, symmetry="vertical", seed=None, noise=None, noise_shift=None):
+    def __init__(
+        self,
+        width=64,
+        height=64,
+        symmetry="vertical",
+        seed=None,
+        noise=None,
+        noise_shift=None,
+    ):
         """Craters
         Builds a craters system of size `width` x `height` and symmetrical across the `symmetry` axis.
 
@@ -135,23 +211,25 @@ def __init__(self, width=64, height=64, symmetry="vertical", seed=None, noise=No
 
         noise = GameMap.noise(seed, noise, noise_shift, symmetry)
 
-        min_craters = max(2, width*height // 1000)
-        max_craters = max(4, width*height // 500)
-        num_craters = noise.random.randint(min_craters, max_craters+1)
+        min_craters = max(2, width * height // 1000)
+        max_craters = max(4, width * height // 500)
+        num_craters = noise.random.randint(min_craters, max_craters + 1)
 
         # Mask = how many craters have hit the spot. When it symmetrizes, it will divide by 2.
         mask = np.zeros((height, width))
         x = np.linspace(0, 1, width)
         y = np.linspace(0, 1, height)
-        crater_noise = noise(x, y, frequency=10) * 0.5 + 0.75 # Don't want perfectly circular craters.
+        crater_noise = (
+            noise(x, y, frequency=10) * 0.5 + 0.75
+        )  # Don't want perfectly circular craters.
         xx, yy = np.mgrid[:width, :height]
 
         # ice should be around edges of crater
         ice_mask = np.zeros((height, width))
         for i in range(num_craters):
             cx, cy = noise.random.randint(width), np.random.randint(height)
-            cr = noise.random.randint(3, min(width, height)//4)
-            c = (xx - cx) **2 + (yy - cy)**2
+            cr = noise.random.randint(3, min(width, height) // 4)
+            c = (xx - cx) ** 2 + (yy - cy) ** 2
             c = c.T * crater_noise
             mask[c < cr**2] += 1
             edge = np.logical_and(c >= cr**2, c < 2 * cr**2)
@@ -162,18 +240,19 @@ def __init__(self, width=64, height=64, symmetry="vertical", seed=None, noise=No
         # symmetrize(mask, symmetry)
         # symmetrize(ice_mask, symmetry)
 
-        rubble = (np.minimum(mask, 1) * 90 + 10) * noise(x, y+100, frequency=3)
-        ice = noise(x, y-100)
-        ice[ice_mask==0] = 0
+        rubble = (np.minimum(mask, 1) * 90 + 10) * noise(x, y + 100, frequency=3)
+        ice = noise(x, y - 100)
+        ice[ice_mask == 0] = 0
         ice[ice < np.percentile(ice, 95)] = 0
         ice = np.round(50 * ice + 50)
 
-        ore = np.minimum(mask, 1) * noise(x, y+100, frequency=3)
+        ore = np.minimum(mask, 1) * noise(x, y + 100, frequency=3)
         ore[ore < np.percentile(ore, 95)] = 0
         ore = np.round(50 * ore + 50)
 
         super().__init__(rubble, ice, ore, symmetry)
 
+
 def solve_poisson(f):
     """
     Solves the Poisson equation
@@ -199,25 +278,76 @@ def solve_poisson(f):
     potential = idctn(dct, type=1)
     return potential / 2
 
+
 def nabla(x):
     return convolve(x, ((0, 1, 0), (1, -5, 1), (0, 1, 0)))
 
+
 def dxx(x):
     return convolve(x, ((1,), (-2,), (1,)))
 
+
 def dyy(x):
     return convolve(x, ((1, -2, 1),))
 
+
 def dxy(x):
     return convolve(x, ((1, 0, -1), (0, 0, 0), (-1, 0, 1))) / 4
 
+
 def laplacian(x):
     return convolve(x, ((1, 1, 1), (1, -8, 1), (1, 1, 1))) / 3
 
 
+@dataclass
+class MountainConfig:
+    # default parameters -> {'rubble_count': 21.35736545138889, 'ice': 36.19, 'ore': 36.0}
+
+    # controls amount of ice spread along mountain peaks
+    ice_high_range: Tuple[float, float] = (98.9, 100)
+    # around middle level
+    ice_mid_range: Tuple[float, float] = (52.5, 53)
+    # around lower level
+    ice_low_range: Tuple[float, float] = (0, 21)
+    # controls amount of ore spread along middle of the way to the mountain peaks
+    ore_mid_range: Tuple[float, float] = (84, 85)
+    # controls amount of ore spread along lower part of the mountain peaks, should be smaller than ore_mid_range
+    ore_low_range: Tuple[float, float] = (61.4, 62)
+
+
+MOUNTAIN_CONFIGS = dict(
+    # {'rubble_count': 21.35736545138889, 'ice': 36.19, 'ore': 36.0}
+    high_ice_high_ore=MountainConfig(),
+    # {'rubble_count': 21.35736545138889, 'ice': 16.02, 'ore': 36.0}
+    low_ice_high_ore=MountainConfig(
+        ice_high_range=(99.5, 100),
+        ice_mid_range=(52.8, 53),
+        ice_low_range=(0, 20),
+    ),
+    # {'rubble_count': 21.35736545138889, 'ice': 36.19, 'ore': 16.0}
+    high_ice_low_ore=MountainConfig(ore_low_range=(61.7, 62), ore_mid_range=(84.6, 85)),
+    # # {'rubble_count': 21.35736545138889, 'ice': 16.02, 'ore': 16.0}
+    low_ice_low_ore=MountainConfig(
+        ice_high_range=(99.5, 100),
+        ice_mid_range=(52.8, 53),
+        ice_low_range=(0, 20),
+        ore_low_range=(61.7, 62),
+        ore_mid_range=(84.6, 85),
+    ),
+)
+
 
 class Mountain(GameMap):
-    def __init__(self, width=64, height=64, symmetry="vertical", seed=None, noise=None, noise_shift=None):
+    def __init__(
+        self,
+        width=64,
+        height=64,
+        symmetry="vertical",
+        seed=None,
+        noise=None,
+        noise_shift=None,
+        config: MountainConfig = MOUNTAIN_CONFIGS["high_ice_high_ore"],
+    ):
         """Mountain
         Builds a mountain system of size `width` x `height` and symmetrical across the `symmetry` axis.
 
@@ -233,9 +363,9 @@ def __init__(self, width=64, height=64, symmetry="vertical", seed=None, noise=No
         f = np.zeros((height, width))
 
         # Sprinkle a few mountains on the map.
-        min_mountains = max(2, width*height//750)
-        max_mountains = max(4, width*height//375)
-        num_mountains = noise.random.randint(min_mountains, max_mountains+1)
+        min_mountains = max(2, width * height // 750)
+        max_mountains = max(4, width * height // 375)
+        num_mountains = noise.random.randint(min_mountains, max_mountains + 1)
         for i in range(num_mountains):
             x, y = noise.random.randint(width), noise.random.randint(height)
             f[y][x] -= 1
@@ -251,16 +381,18 @@ def __init__(self, width=64, height=64, symmetry="vertical", seed=None, noise=No
         mask -= np.amin(mask)
 
         # Find the valleys
-        Lap = convolve(mask, ((1, 1, 1), (1, -8, 1), (1, 1, 1))) / 3 # Laplacian
+        Lap = convolve(mask, ((1, 1, 1), (1, -8, 1), (1, 1, 1))) / 3  # Laplacian
         Dxx = convolve(mask, ((1,), (-2,), (1,)))
         Dyy = convolve(mask, ((1, -2, 1),))
         Dxy = convolve(mask, ((1, 0, -1), (0, 0, 0), (-1, 0, 1)))
-        det = 16 * Dxx * Dyy - Dxy ** 2 # Hessian determinant
-        det = det.astype(np.csingle) # complex
-
-        cond = Lap * (2 * Lap - np.sqrt(4 * Lap**2 - det)) / det - 0.25 # ratio of eigenvalues
-        cond = abs(cond) # should already be real except for floating point errors
-        cond = np.maximum(cond, 1/cond)
+        det = 16 * Dxx * Dyy - Dxy**2  # Hessian determinant
+        det = det.astype(np.csingle)  # complex
+
+        cond = (
+            Lap * (2 * Lap - np.sqrt(4 * Lap**2 - det)) / det - 0.25
+        )  # ratio of eigenvalues
+        cond = abs(cond)  # should already be real except for floating point errors
+        cond = np.maximum(cond, 1 / cond)
         # symmetrize(cond, symmetry) # for floating point errors
 
         def bdry(x, y):
@@ -268,6 +400,7 @@ def bdry(x, y):
             return abs(cond[y][x]) > 20 and f[y][x] == 0
 
         closed_set = set()
+
         def flood_fill(x, y):
             nonlocal closed_set
             region = []
@@ -312,30 +445,50 @@ def flood_fill(x, y):
 
         for x in range(width):
             for y in range(height):
-                if bdry(x, y): mask[y][x] = 0
+                if bdry(x, y):
+                    mask[y][x] = 0
 
-        mask[np.where(mask > 0)] -= np.amin(mask[np.where(mask>0)])
+        mask[np.where(mask > 0)] -= np.amin(mask[np.where(mask > 0)])
         mask -= np.amin(mask)
         mask /= np.amax(mask)
 
-        rubble = (100*mask).round()
-        ice = (100 * mask).round()
-        mid_mask = (ice < np.percentile(ice, 50)) & (ice > np.percentile(ice, 48)) | (ice < np.percentile(ice, 20)) & (ice > np.percentile(ice, 0))
-        ice[ice < np.percentile(ice, 98)] = 0
-        ice[ice != 0] = 1
-        ice[mid_mask] = 1
-        
-        ore = (100 * mask).round()
-        mid_mask = (ore < np.percentile(ore, 62)) & (ore > np.percentile(ore, 60))
-        ore[ore < np.percentile(ore, 83.5)] = 0
-        ore[ore > np.percentile(ore, 84.5)] = 0
-        ore[ore != 0] = 1
-        ore[mid_mask] = 1
+        rubble = (100 * mask).round()
+        ice = 100 * mask
+        high_mask = (ice > np.percentile(ice, config.ice_high_range[0])) & (
+            ice < np.percentile(ice, config.ice_high_range[1])
+        )
+        mid_mask = (ice > np.percentile(ice, config.ice_mid_range[0])) & (
+            ice < np.percentile(ice, config.ice_mid_range[1])
+        )
+        low_mask = (ice > np.percentile(ice, config.ice_low_range[0])) & (
+            ice < np.percentile(ice, config.ice_low_range[1])
+        )
+        ice = low_mask | mid_mask | high_mask
+
+        ore = 100 * mask
+
+        mid_mask = (ore > np.percentile(ore, config.ore_mid_range[0])) & (
+            ore < np.percentile(ore, config.ore_mid_range[1])
+        )
+        low_mask = (ore > np.percentile(ore, config.ore_low_range[0])) & (
+            ore < np.percentile(ore, config.ore_low_range[1])
+        )
+
+        ore = low_mask | mid_mask
 
         super().__init__(rubble, ice, ore, symmetry)
 
+
 class Island(GameMap):
-    def __init__(self, width=64, height=64, symmetry="vertical", seed=None, noise=None, noise_shift=None):
+    def __init__(
+        self,
+        width=64,
+        height=64,
+        symmetry="vertical",
+        seed=None,
+        noise=None,
+        noise_shift=None,
+    ):
         """Island
         Builds an island system of size `width` x `height` and symmetrical across the `symmetry` axis.
 
@@ -353,9 +506,9 @@ def __init__(self, width=64, height=64, symmetry="vertical", seed=None, noise=No
         mask[mask >= 1] = 1
 
         s = -1
-        while np.sum(mask==0) != s:
-            s = np.sum(mask==0)
-            mask = convolve(mask, [[1]*3]*3, mode="constant", cval=1) // 6
+        while np.sum(mask == 0) != s:
+            s = np.sum(mask == 0)
+            mask = convolve(mask, [[1] * 3] * 3, mode="constant", cval=1) // 6
 
         # Shift every spot on the map a little, making the islands nosier.
         x = np.linspace(0, 1, width)
@@ -365,36 +518,46 @@ def __init__(self, width=64, height=64, symmetry="vertical", seed=None, noise=No
 
         xx, yy = np.mgrid[:width, :height]
         new_xx = xx + np.round(noise_dx).astype(int)
-        new_xx = np.maximum(0, np.minimum(width-1, new_xx))
+        new_xx = np.maximum(0, np.minimum(width - 1, new_xx))
         new_yy = yy + np.round(noise_dy).astype(int)
-        new_yy = np.maximum(0, np.minimum(height-1, new_yy))
+        new_yy = np.maximum(0, np.minimum(height - 1, new_yy))
 
         mask[yy, xx] = mask[new_yy, new_xx]
 
         # symmetrize(mask, symmetry)
 
         rubble = noise(x, y - 100, frequency=3) * 50 + 50
-        rubble[mask==0] //= 20
+        rubble[mask == 0] //= 20
 
         # Unsure what to do about ice, ore right now. Place in pockets on islands?
         ice = noise(x, y + 200, frequency=10) ** 2 * 100
         ice[ice < 50] = 0
-        ice[mask!=0] = 0
+        ice[mask != 0] = 0
 
         ore = noise(x, y - 200, frequency=10) ** 2 * 100
         ore[ore < 50] = 0
-        ore[mask!=0] = 0
+        ore[mask != 0] = 0
 
         super().__init__(rubble, ice, ore, symmetry)
 
 
 if __name__ == "__main__":
     import argparse
+
     parser = argparse.ArgumentParser(description="Generate maps for Lux AI 2022.")
-    parser.add_argument("-t", "--map_type", help="Map type ('Cave', 'Craters', 'Island', 'Mountain')", required=False)
+    parser.add_argument(
+        "-t",
+        "--map_type",
+        help="Map type ('Cave', 'Craters', 'Island', 'Mountain')",
+        required=False,
+    )
     parser.add_argument("-s", "--size", help="Size (32-64)", required=False)
     parser.add_argument("-d", "--seed", help="Seed")
-    parser.add_argument("-m", "--symmetry", help="Symmetry ('horizontal', 'rotational', 'vertical', '/', '\\')")
+    parser.add_argument(
+        "-m",
+        "--symmetry",
+        help="Symmetry ('horizontal', 'rotational', 'vertical', '/', '\\')",
+    )
 
     args = vars(parser.parse_args())
     map_type = args.get("map_type", None)
@@ -407,6 +570,8 @@ def __init__(self, width=64, height=64, symmetry="vertical", seed=None, noise=No
         seed = int(args["seed"])
     else:
         seed = None
-    game_map = GameMap.random_map(seed=seed, symmetry=symmetry, map_type=map_type, width=width, height=height)
+    game_map = GameMap.random_map(
+        seed=seed, symmetry=symmetry, map_type=map_type, width=width, height=height
+    )
     viz(game_map)
     input()
diff --git a/luxai_s2/luxai_s2/map_generator/symnoise.py b/luxai_s2/luxai_s2/map_generator/symnoise.py
index 50aeaf4a..3545bee4 100644
--- a/luxai_s2/luxai_s2/map_generator/symnoise.py
+++ b/luxai_s2/luxai_s2/map_generator/symnoise.py
@@ -1,23 +1,25 @@
-from vec_noise import snoise2
 from collections.abc import Iterable
+
 import numpy as np
+from vec_noise import snoise2
+
 
 def symmetrize(x, symmetry="vertical"):
     # In place operation to average along the symmetry.
     height, width = x.shape
     if symmetry == "horizontal":
-        x[height//2:] += x[(height-1)//2::-1]
-        x[(height-1)//2::-1] = x[height//2:]
+        x[height // 2 :] += x[(height - 1) // 2 :: -1]
+        x[(height - 1) // 2 :: -1] = x[height // 2 :]
     elif symmetry == "vertical":
-        x[:, width//2:] += x[:, (width-1)//2::-1]
-        x[:, (width-1)//2::-1] = x[:, width//2:]
+        x[:, width // 2 :] += x[:, (width - 1) // 2 :: -1]
+        x[:, (width - 1) // 2 :: -1] = x[:, width // 2 :]
     elif symmetry == "rotational":
-        x[height//2-1:, :] += x[(height+1)//2::-1, ::-1]
-        x[(height+1)//2::-1, ::-1] = x[height//2-1:, :]
+        x[height // 2 - 1 :, :] += x[(height + 1) // 2 :: -1, ::-1]
+        x[(height + 1) // 2 :: -1, ::-1] = x[height // 2 - 1 :, :]
     elif symmetry == "/":
         for j in range(height):
-            x[j, -j-1::-1] += x[j:, -j-1]
-            x[j:, -j-1] = x[j, -j-1::-1]
+            x[j, -j - 1 :: -1] += x[j:, -j - 1]
+            x[j:, -j - 1] = x[j, -j - 1 :: -1]
     elif symmetry == "\\":
         for j in range(height):
             x[j, j:] += x[j:, j]
@@ -25,13 +27,22 @@ def symmetrize(x, symmetry="vertical"):
     else:
         x *= 2
 
-    if x.dtype.kind == "i": # Integer arrays need integer division.
+    if x.dtype.kind == "i":  # Integer arrays need integer division.
         x //= 2
     else:
         x /= 2
 
+
 class SymmetricNoise(object):
-    def __init__(self, seed: int = 0, octaves: int = 1, symmetry="vertical", width=None, height=None, noise_shift=0):
+    def __init__(
+        self,
+        seed: int = 0,
+        octaves: int = 1,
+        symmetry="vertical",
+        width=None,
+        height=None,
+        noise_shift=0,
+    ):
         """Symmetrical Simplex noise.
 
             ex.: noise = SymmetricalNoise(symmetry="rotational", width=50, height=100,
@@ -45,7 +56,9 @@ def __init__(self, seed: int = 0, octaves: int = 1, symmetry="vertical", width=N
             octaves : how fine the features are
         """
         if symmetry not in (None, "vertical", "horizontal", "rotational", "/", "\\"):
-            raise ValueError("symmetry must be one of None, 'vertical', 'horizontal', 'rotational', '/', and '\\'")
+            raise ValueError(
+                "symmetry must be one of None, 'vertical', 'horizontal', 'rotational', '/', and '\\'"
+            )
         if symmetry and symmetry in "/\\" and width != height:
             raise ValueError("width and height must be equal if symmetry = / or \\")
 
@@ -63,7 +76,7 @@ def __init__(self, seed: int = 0, octaves: int = 1, symmetry="vertical", width=N
     def update_symmetry(self, symmetry):
         self.symmetry = symmetry
 
-    def noise(self, x=None, y=None, frequency: float=1):
+    def noise(self, x=None, y=None, frequency: float = 1):
         # x and y can be arrays, but all values should be between 0 and 1.
         if x is None and self.width is None:
             raise ValueError("Must provide x or define width in initialization")
@@ -88,10 +101,11 @@ def noise(self, x=None, y=None, frequency: float=1):
     def __call__(self, *args, **kwargs):
         return self.noise(*args, **kwargs)
 
+
 if __name__ == "__main__":
     import matplotlib.pyplot as plt
-    from mpl_toolkits.axes_grid1 import ImageGrid
     import numpy as np
+    from mpl_toolkits.axes_grid1 import ImageGrid
 
     SEED = 0
     w, h = 100, 100
@@ -102,11 +116,13 @@ def __call__(self, *args, **kwargs):
         noise = SymmetricNoise(symmetry=s, octaves=3, seed=SEED)
         imgs.append(noise(x, y))
 
-    fig = plt.figure(figsize=(10., 2.))
-    grid = ImageGrid(fig, 111,
-                     nrows_ncols=(1, 5),
-                     axes_pad=0.1,
-                     )
+    fig = plt.figure(figsize=(10.0, 2.0))
+    grid = ImageGrid(
+        fig,
+        111,
+        nrows_ncols=(1, 5),
+        axes_pad=0.1,
+    )
 
     for ax, im in zip(grid, imgs):
         # Iterating over the grid returns the Axes.
diff --git a/luxai_s2/luxai_s2/map_generator/timing.py b/luxai_s2/luxai_s2/map_generator/timing.py
index 1bc793e8..f09a2a68 100644
--- a/luxai_s2/luxai_s2/map_generator/timing.py
+++ b/luxai_s2/luxai_s2/map_generator/timing.py
@@ -2,7 +2,9 @@
 
 N = 100
 for map_type in ["Cave", "Craters", "Island", "Mountain"]:
-    time = timeit.timeit(f"{map_type}(64, 64)", setup=f"from generator import {map_type}", number=N)
+    time = timeit.timeit(
+        f"{map_type}(64, 64)", setup=f"from generator import {map_type}", number=N
+    )
     print(f"Took {time / N * 1000:.2f}ms for {map_type}.")
 
 
diff --git a/luxai_s2/luxai_s2/map_generator/visualize.py b/luxai_s2/luxai_s2/map_generator/visualize.py
index cdfbf3c9..b174a3ed 100644
--- a/luxai_s2/luxai_s2/map_generator/visualize.py
+++ b/luxai_s2/luxai_s2/map_generator/visualize.py
@@ -9,14 +9,14 @@
 def viz(game_map, screen=None):
     N = min(500 // game_map.width, 750 // game_map.height)
     if screen is None:
-        screen = pygame.display.set_mode((3*N*game_map.width, N*game_map.height))
+        screen = pygame.display.set_mode((3 * N * game_map.width, N * game_map.height))
 
     for x in range(game_map.width):
         for y in range(game_map.height):
-            rubble_color = [255-game_map.rubble[y][x]*255/100]*3
-            ice_color = [0, 0, game_map.ice[y][x]*255/100]
-            ore_color = [game_map.ore[y][x]*255/100, 0, 0]
-            screen.fill(rubble_color, (N*x, N*y, N, N))
-            screen.fill(ice_color, (N*x+N*game_map.width, N*y, N, N))
-            screen.fill(ore_color, (N*x+2*N*game_map.width, N*y, N, N))
+            rubble_color = [255 - game_map.rubble[y][x] * 255 / 100] * 3
+            ice_color = [0, 0, game_map.ice[y][x] * 255 / 100]
+            ore_color = [game_map.ore[y][x] * 255 / 100, 0, 0]
+            screen.fill(rubble_color, (N * x, N * y, N, N))
+            screen.fill(ice_color, (N * x + N * game_map.width, N * y, N, N))
+            screen.fill(ore_color, (N * x + 2 * N * game_map.width, N * y, N, N))
     pygame.display.update()
diff --git a/luxai_s2/luxai_s2/map_generator/visualize_samples.py b/luxai_s2/luxai_s2/map_generator/visualize_samples.py
index 1dd34d18..a8b53b38 100644
--- a/luxai_s2/luxai_s2/map_generator/visualize_samples.py
+++ b/luxai_s2/luxai_s2/map_generator/visualize_samples.py
@@ -1,27 +1,27 @@
 from math import prod
 from pathlib import Path
-from luxai_s2.env import LuxAI_S2
-import pygame
+
 import numpy as np
+import pygame
+
+from luxai_s2.env import LuxAI_S2
+
 if __name__ == "__main__":
     env = LuxAI_S2()
     N = 100
     Path("./map_samples").mkdir(parents=True, exist_ok=True)
-    stats = dict(
-        rubble_count=[],
-        ice=[],
-        ore=[]
-    )
-    for seed in range(10000, 10000+N):
+    stats = dict(rubble_count=[], ice=[], ore=[])
+    for seed in range(10000, 10000 + N):
         env.reset(seed)
         env.render()
         env.py_visualizer.update_scene(env.state)
-        pygame.image.save(env.py_visualizer.screen , f"map_samples/{seed}.jpg")
+
+        pygame.image.save(env.py_visualizer.screen, f"map_samples/{seed}.jpg")
         stats["ice"] += [env.state.board.ice.sum()]
         stats["ore"] += [env.state.board.ore.sum()]
         stats["rubble_count"] += [env.state.board.rubble.sum()]
 
-    for k in stats: 
+    for k in stats:
         stats[k] = np.mean(stats[k])
     stats["rubble_count"] /= prod(env.state.board.rubble.shape)
     print(stats)
diff --git a/luxai_s2/luxai_s2/pyvisual/visualizer.py b/luxai_s2/luxai_s2/pyvisual/visualizer.py
index 0dc16574..215af615 100644
--- a/luxai_s2/luxai_s2/pyvisual/visualizer.py
+++ b/luxai_s2/luxai_s2/pyvisual/visualizer.py
@@ -1,4 +1,5 @@
 import numpy as np
+
 try:
     import pygame
     from pygame import gfxdraw
@@ -7,19 +8,34 @@
 from luxai_s2.map.board import Board
 from luxai_s2.state import State
 from luxai_s2.unit import UnitType
+
 try:
     import matplotlib.pyplot as plt
-    color_to_rgb = dict(yellow=[236, 238, 126], green=[173, 214, 113], blue=[154, 210, 203], red=[164, 74, 63])
-    strain_colors=plt.colormaps['Pastel1']
+
+    color_to_rgb = dict(
+        yellow=[236, 238, 126],
+        green=[173, 214, 113],
+        blue=[154, 210, 203],
+        red=[164, 74, 63],
+    )
+    strain_colors = plt.colormaps["Pastel1"]
 except:
     pass
+
+
 class Visualizer:
     def __init__(self, state: State) -> None:
         # self.screen = pygame.display.set_mode((3*N*game_map.width, N*game_map.height))
         self.screen_size = (1000, 1000)
         self.board = state.board
-        self.tile_width = min(self.screen_size[0] // self.board.width, self.screen_size[1] // self.board.height)
-        self.WINDOW_SIZE = (self.tile_width * self.board.width, self.tile_width * self.board.height)
+        self.tile_width = min(
+            self.screen_size[0] // self.board.width,
+            self.screen_size[1] // self.board.height,
+        )
+        self.WINDOW_SIZE = (
+            self.tile_width * self.board.width,
+            self.tile_width * self.board.height,
+        )
         self.surf = pygame.Surface(self.WINDOW_SIZE)
         self.surf.fill([239, 120, 79])
         self.state = state
@@ -34,45 +50,76 @@ def init_window(self):
     def rubble_color(self, rubble):
         opacity = 0.2 + min(rubble / 100, 1) * 0.8
         return [96, 32, 9, opacity * 255]
+
     def ore_color(self, rubble):
         return [218, 167, 48, 255]
+
     def ice_color(self, rubble):
         return [44, 158, 211, 255]
 
-
     def update_scene(self, state: State):
         self.state = state
         self.surf.fill([239, 120, 79, 255])
         for x in range(self.board.width):
             for y in range(self.board.height):
                 rubble_amt = self.state.board.rubble[x][y]
-                rubble_color = self.rubble_color(rubble_amt) #[255 - self.state.board.rubble[y][x] * 255 / 100] * 3
+                rubble_color = self.rubble_color(
+                    rubble_amt
+                )  # [255 - self.state.board.rubble[y][x] * 255 / 100] * 3
                 # import ipdb;ipdb.set_trace()
-                gfxdraw.box(self.surf, (self.tile_width * x, self.tile_width * y, self.tile_width, self.tile_width), rubble_color)
+                gfxdraw.box(
+                    self.surf,
+                    (
+                        self.tile_width * x,
+                        self.tile_width * y,
+                        self.tile_width,
+                        self.tile_width,
+                    ),
+                    rubble_color,
+                )
                 if self.state.board.ice[x, y] > 0:
                     pygame.draw.rect(
                         self.surf,
                         self.ice_color(rubble_amt),
-                        pygame.Rect(self.tile_width * x, self.tile_width * y, self.tile_width, self.tile_width),
+                        pygame.Rect(
+                            self.tile_width * x,
+                            self.tile_width * y,
+                            self.tile_width,
+                            self.tile_width,
+                        ),
                     )
                 # if self.state.valid_spawns_mask
                 if self.state.board.ore[x, y] > 0:
                     pygame.draw.rect(
                         self.surf,
                         self.ore_color(rubble_amt),
-                        pygame.Rect(self.tile_width * x, self.tile_width * y, self.tile_width, self.tile_width),
+                        pygame.Rect(
+                            self.tile_width * x,
+                            self.tile_width * y,
+                            self.tile_width,
+                            self.tile_width,
+                        ),
                     )
                 if self.state.board.lichen_strains[x, y] != -1:
-                    c = strain_colors.colors[self.state.board.lichen_strains[x, y] % len(strain_colors.colors)]
+                    c = strain_colors.colors[
+                        self.state.board.lichen_strains[x, y]
+                        % len(strain_colors.colors)
+                    ]
                     pygame.draw.rect(
                         self.surf,
-                        [int(c[0]*255),int(c[1]*255),int(c[2]*255)],
-                        pygame.Rect(self.tile_width * x, self.tile_width * y, self.tile_width, self.tile_width),
+                        [int(c[0] * 255), int(c[1] * 255), int(c[2] * 255)],
+                        pygame.Rect(
+                            self.tile_width * x,
+                            self.tile_width * y,
+                            self.tile_width,
+                            self.tile_width,
+                        ),
                     )
         if len(state.teams) > 0:
-            team_color = dict(player_0=[224,65,40], player_1=[0, 112, 81])
+            team_color = dict(player_0=[224, 65, 40], player_1=[0, 112, 81])
             for agent in state.factories:
-                if agent not in state.teams: continue
+                if agent not in state.teams:
+                    continue
                 team = state.teams[agent]
                 for factory in state.factories[agent].values():
                     x = factory.pos.x
@@ -87,21 +134,24 @@ def update_scene(self, state: State):
                             self.tile_width * 3,
                             self.tile_width * 3,
                         ),
-                        border_radius=int(self.tile_width / 2)
+                        border_radius=int(self.tile_width / 2),
+                    )
+                    self.sans_font = pygame.font.SysFont("Open Sans", 30)
+                    self.surf.blit(
+                        self.sans_font.render("F", False, [51, 56, 68]),
+                        (self.tile_width * x, self.tile_width * y),
                     )
-                    self.sans_font = pygame.font.SysFont('Open Sans', 30)
-                    self.surf.blit(self.sans_font.render('F', False, [51,56,68]), (self.tile_width * x, self.tile_width * y))
             for agent in state.units:
-                if agent not in state.teams: continue
+                if agent not in state.teams:
+                    continue
                 team = state.teams[agent]
                 for unit in state.units[agent].values():
                     x = unit.pos.x
                     y = unit.pos.y
-                    h=1
+                    h = 1
                     pygame.draw.rect(
                         self.surf,
-                        [51,56,68],
-                        
+                        [51, 56, 68],
                         pygame.Rect(
                             self.tile_width * (x),
                             self.tile_width * (y),
@@ -114,21 +164,26 @@ def update_scene(self, state: State):
                         # color_to_rgb[team.faction.value.color],
                         team_color[agent],
                         pygame.Rect(
-                            self.tile_width * (x)+h,
-                            self.tile_width * (y)+ h,
+                            self.tile_width * (x) + h,
+                            self.tile_width * (y) + h,
                             (self.tile_width) * 1 - h * 2,
                             (self.tile_width) * 1 - h * 2,
                         ),
                     )
-                    
+
                     label = "H"
                     if unit.unit_type == UnitType.LIGHT:
                         label = "L"
-                    self.sans_font = pygame.font.SysFont('Open Sans', 20)
-                    self.surf.blit(self.sans_font.render(label, False, [51,56,68]), (self.tile_width * x+2, self.tile_width * y+2))
+                    self.sans_font = pygame.font.SysFont("Open Sans", 20)
+                    self.surf.blit(
+                        self.sans_font.render(label, False, [51, 56, 68]),
+                        (self.tile_width * x + 2, self.tile_width * y + 2),
+                    )
+
     def render(self):
         pygame.display.update()
         self.screen.blit(self.surf, (0, 0))
+
     def _create_image_array(self, screen, size):
         scaled_screen = pygame.transform.smoothscale(screen, size)
         return np.transpose(
diff --git a/luxai_s2/luxai_s2/replay/__init__.py b/luxai_s2/luxai_s2/replay/__init__.py
index ef085fb8..feb6bff8 100644
--- a/luxai_s2/luxai_s2/replay/__init__.py
+++ b/luxai_s2/luxai_s2/replay/__init__.py
@@ -1,1 +1,1 @@
-from .replay import decode_replay_file, generate_replay
\ No newline at end of file
+from .replay import decode_replay_file, generate_replay
diff --git a/luxai_s2/luxai_s2/replay/replay.py b/luxai_s2/luxai_s2/replay/replay.py
index cdca3c31..e4f6f480 100644
--- a/luxai_s2/luxai_s2/replay/replay.py
+++ b/luxai_s2/luxai_s2/replay/replay.py
@@ -1,8 +1,11 @@
+import json
 import os
+import os.path as osp
 from typing import List, Tuple
+
 from luxai_s2.state import State
-import os.path as osp
-import json
+
+
 def decode_replay_file(replay_file) -> Tuple[State, List]:
     """
     Takes an input replay file of any kind from any source and extracts the full trajectory with observations
@@ -18,9 +21,10 @@ def decode_replay_file(replay_file) -> Tuple[State, List]:
     elif ext == ".h5":
         pass
 
+
 def generate_replay(states: List[State]):
     """
     Generates a compressed replay.
-    
+
     """
-    return [s.get_compressed_obs() for s in states]
\ No newline at end of file
+    return [s.get_compressed_obs() for s in states]
diff --git a/luxai_s2/luxai_s2/spaces/act_space.py b/luxai_s2/luxai_s2/spaces/act_space.py
index 5b8667ac..dba96490 100644
--- a/luxai_s2/luxai_s2/spaces/act_space.py
+++ b/luxai_s2/luxai_s2/spaces/act_space.py
@@ -1,18 +1,20 @@
+import random
 from typing import Any, Dict, List
-from gym import spaces
+
 import gym
+import numpy as np
+from gym import spaces
 
 from luxai_s2.config import EnvConfig
 from luxai_s2.factory import Factory
 from luxai_s2.team import FactionTypes
 from luxai_s2.unit import Unit
-import random
-import numpy as np
+
 
 class FactionString(gym.Space):
     def __init__(
-                self,
-            ):
+        self,
+    ):
         self.valid_factions = ["Null"] + [x.name for x in FactionTypes]
 
     def sample(self):
@@ -21,6 +23,7 @@ def sample(self):
     def contains(self, x):
         return type(x) == "str" and x in self.valid_factions
 
+
 class PartialDict(spaces.Dict):
     def contains(self, x: Any) -> bool:
         if not isinstance(x, dict) or len(x) > len(self.spaces):
@@ -34,26 +37,31 @@ def contains(self, x: Any) -> bool:
                 return False, f"{k} is not on your team or does not exist"
         return True, None
 
+
 class ActionsQueue(spaces.Space):
     def __init__(self, action_space: spaces.Space, max_length: int) -> None:
         super().__init__((), float)
         self.max_length = max_length
         self.action_space = action_space
+
     def sample(self):
         queue_size = self.np_random.randint(0, self.max_length) + 1
         action_q = []
         for _ in range(queue_size):
             action_q.append(self.action_space.sample())
         return action_q
+
     def contains(self, x: Any) -> bool:
         if isinstance(x, list) and len(x) > self.max_length:
             return False
-        elif isinstance(x, np.ndarray) and len(x.shape) == 2 and len(x) > self.max_length:
+        elif (
+            isinstance(x, np.ndarray) and len(x.shape) == 2 and len(x) > self.max_length
+        ):
             return False
         elif isinstance(x, np.ndarray) and len(x.shape) == 1:
             x = [x]
         # if (not isinstance(x, list) and not isinstance(x, np.ndarray)) or len(x) > self.max_length:
-            # return False
+        # return False
         for a in x:
             # fix issue where multidiscrete space does not do 100% check on type of value
             try:
@@ -63,18 +71,24 @@ def contains(self, x: Any) -> bool:
                 return False
         return True
 
+
 def get_act_space_init(config: EnvConfig, agent: str):
     # Get action space for turn 0 initialization
     act_space = dict()
     act_space["faction"] = FactionString()
-    act_space["spawns"] = spaces.Box(low=0, high=config.map_size, shape=(config.MAX_FACTORIES, 2), dtype=int)
+    act_space["spawns"] = spaces.Box(
+        low=0, high=config.map_size, shape=(config.MAX_FACTORIES, 2), dtype=int
+    )
     return spaces.Dict(act_space)
 
+
 def get_act_space_bid(config: EnvConfig, agent: str):
     act_space = dict()
     act_space["faction"] = FactionString()
     act_space["bid"] = spaces.Discrete(100000)
     return spaces.Dict(act_space)
+
+
 def get_act_space_placement(config: EnvConfig, agent: str):
     # Get action space for turn 0 initialization
     act_space = dict()
@@ -83,7 +97,13 @@ def get_act_space_placement(config: EnvConfig, agent: str):
     act_space["metal"] = spaces.Discrete(100000)
     return spaces.Dict(act_space)
 
-def get_act_space(units: Dict[str, Dict[str, Unit]], factories: Dict[str, Dict[str, Factory]], config: EnvConfig, agent: str):
+
+def get_act_space(
+    units: Dict[str, Dict[str, Unit]],
+    factories: Dict[str, Dict[str, Factory]],
+    config: EnvConfig,
+    agent: str,
+):
     act_space = dict()
 
     # for consistency, every action space per unit is fixed, makes it easier to work out of the box.
@@ -109,14 +129,21 @@ def get_act_space(units: Dict[str, Dict[str, Unit]], factories: Dict[str, Dict[s
         # a[3] = X, amount of resources transferred or picked up if action is transfer or pickup.
         # If action is recharge, it is how much energy to store before executing the next action in queue
 
-        # a[4] = -1,0,1,... - -1 means repeat infinitely after succesful execution. 0 means never repeat again. n > 0 means repeat n times.
+        # a[4] = 0 or 1 (false or true). If True, then action is placed to the back of the action queue after it has been exhausted
+
+        # a[5] = X, number of times to execute this action before exhausting it and removing it from the front of the action queue. Minimum is 1.
         act_space[u.unit_id] = ActionsQueue(
-            spaces.Box(low=np.array([0,0,0,0,-1]), high=np.array([5,4,4,config.max_transfer_amount, 100000]), shape=(5,), dtype=np.int64),
-        config.UNIT_ACTION_QUEUE_SIZE)
+            spaces.Box(
+                low=np.array([0, 0, 0, 0, 0, 1]),
+                high=np.array([5, 4, 4, config.max_transfer_amount, 1, 9999]),
+                shape=(6,),
+                dtype=np.int64,
+            ),
+            config.UNIT_ACTION_QUEUE_SIZE,
+        )
 
     for factory in factories[agent].values():
         # action type (0 = build light robot, 1 = build heavy robot, 2 = water lichen)
         act_space[factory.unit_id] = spaces.Discrete(3)
 
-
     return PartialDict(act_space)
diff --git a/luxai_s2/luxai_s2/spaces/obs_space.py b/luxai_s2/luxai_s2/spaces/obs_space.py
index ef002d45..5188bad1 100644
--- a/luxai_s2/luxai_s2/spaces/obs_space.py
+++ b/luxai_s2/luxai_s2/spaces/obs_space.py
@@ -1,30 +1,37 @@
 import random
 from typing import Any, List
+
+import numpy as np
 from gym import spaces
 
 from luxai_s2.config import EnvConfig
 from luxai_s2.spaces.act_space import ActionsQueue, FactionString
-import numpy as np
+
+
 class DynamicArray(spaces.Space):
     def __init__(self, element_space: spaces.Space, max_length: int) -> None:
         super().__init__((), float)
         self.max_length = max_length
         self.element_space = element_space
+
     def sample(self):
         queue_size = self.np_random.randint(0, self.max_length) + 1
         arr = []
         for _ in range(queue_size):
             arr.append(self.element_space.sample())
         return arr
+
     def contains(self, x: Any) -> bool:
         if isinstance(x, list) and len(x) > self.max_length:
             return False
-        elif isinstance(x, np.ndarray) and len(x.shape) == 2 and len(x) > self.max_length:
+        elif (
+            isinstance(x, np.ndarray) and len(x.shape) == 2 and len(x) > self.max_length
+        ):
             return False
         elif isinstance(x, np.ndarray) and len(x.shape) == 1:
             x = [x]
         # if (not isinstance(x, list) and not isinstance(x, np.ndarray)) or len(x) > self.max_length:
-            # return False
+        # return False
         for a in x:
             # fix issue where multidiscrete space does not do 100% check on type of value
             try:
@@ -33,50 +40,52 @@ def contains(self, x: Any) -> bool:
             except:
                 return False
         return True
+
+
 class UnitTypeSpace(spaces.Space):
     def __init__(
-                self,
-            ):
+        self,
+    ):
         pass
+
     def sample(self):
         return random.choice(["LIGHT", "HEAVY"])
 
     def contains(self, x):
         return type(x) == "str" and (x == "LIGHT" or x == "HEAVY")
+
+
 class FactoryIDSpace(spaces.Space):
     def __init__(
-                self,
-            ):
+        self,
+    ):
         pass
+
     def sample(self):
         return "factory_x"
 
     def contains(self, x):
         return type(x) == "str" and "factory" in x
+
+
 class UnitIDSpace(spaces.Space):
     def __init__(
-                self,
-            ):
+        self,
+    ):
         pass
+
     def sample(self):
         return "unit_x"
 
     def contains(self, x):
         return type(x) == "str" and "unit" in x
 
+
 def get_obs_space(config: EnvConfig, agent_names: List[str], agent: int = 0):
-    max_space = 2 ** 32 - 1
+    max_space = 2**32 - 1
     obs_space = dict()
-    
-    obs_space["real_env_steps"] = spaces.Discrete(config.max_episode_length)
 
-    # weather schedule
-    obs_space["weather_schedule"] = spaces.Box(
-        low=0,
-        high=len(config.WEATHER),
-        shape=(config.max_episode_length,),
-        dtype=int,
-    )
+    obs_space["real_env_steps"] = spaces.Discrete(config.max_episode_length)
 
     # teams obs space
     teams_obs_space = dict()
@@ -87,8 +96,10 @@ def get_obs_space(config: EnvConfig, agent_names: List[str], agent: int = 0):
             water=spaces.Discrete(max_space),
             metal=spaces.Discrete(max_space),
             factories_to_place=spaces.Discrete(config.MAX_FACTORIES),
-            factory_strains=DynamicArray(spaces.Discrete(max_space), max_length=(config.MAX_FACTORIES + 1)),
-            place_first=spaces.Discrete(2)
+            factory_strains=DynamicArray(
+                spaces.Discrete(max_space), max_length=(config.MAX_FACTORIES + 1)
+            ),
+            place_first=spaces.Discrete(2),
         )
     obs_space["teams"] = spaces.Dict(teams_obs_space)
 
@@ -101,7 +112,7 @@ def get_obs_space(config: EnvConfig, agent_names: List[str], agent: int = 0):
         lichen=spaces.Box(low=0, high=99999, shape=map_shape, dtype=int),
         lichen_strains=spaces.Box(low=-1, high=99999, shape=map_shape, dtype=int),
         valid_spawns_mask=spaces.Box(low=0, high=1, shape=map_shape, dtype=int),
-        factories_per_team=spaces.Discrete(config.MAX_FACTORIES)
+        factories_per_team=spaces.Discrete(config.MAX_FACTORIES),
     )
     obs_space["board"] = spaces.Dict(board_obs_space)
 
@@ -125,7 +136,14 @@ def get_obs_space(config: EnvConfig, agent_names: List[str], agent: int = 0):
         )
         if config.UNIT_ACTION_QUEUE_SIZE != 1:
             # note, action queue space copied over from act_space.py
-            obs_dict["action_queue"] = ActionsQueue(spaces.Box(low=np.array([0,0,0,0,-1]), high=np.array([5,4,4,config.max_transfer_amount, 100000]), shape=(5,), dtype=np.int64))
+            obs_dict["action_queue"] = ActionsQueue(
+                spaces.Box(
+                    low=np.array([0, 0, 0, 0, -1]),
+                    high=np.array([5, 4, 4, config.max_transfer_amount, 100000]),
+                    shape=(5,),
+                    dtype=np.int64,
+                )
+            )
         units_obs_space[agent_name] = spaces.Dict(obs_dict)
 
     obs_space["units"] = spaces.Dict(units_obs_space)
@@ -146,7 +164,7 @@ def get_obs_space(config: EnvConfig, agent_names: List[str], agent: int = 0):
             ),
             unit_id=FactoryIDSpace(),
             team_id=spaces.Discrete(2),
-            strain_id=spaces.Discrete(max_space)
+            strain_id=spaces.Discrete(max_space),
         )
         factories_obs_space[agent_name] = spaces.Dict(obs_dict)
     obs_space["factories"] = spaces.Dict(factories_obs_space)
diff --git a/luxai_s2/luxai_s2/state/__init__.py b/luxai_s2/luxai_s2/state/__init__.py
index 2fc6ae1d..3560a068 100644
--- a/luxai_s2/luxai_s2/state/__init__.py
+++ b/luxai_s2/luxai_s2/state/__init__.py
@@ -1,1 +1,1 @@
-from .state import State
\ No newline at end of file
+from .state import State
diff --git a/luxai_s2/luxai_s2/state/state.py b/luxai_s2/luxai_s2/state/state.py
index f30bd69b..b380cfca 100644
--- a/luxai_s2/luxai_s2/state/state.py
+++ b/luxai_s2/luxai_s2/state/state.py
@@ -1,15 +1,16 @@
+import copy
+from collections import OrderedDict
 from dataclasses import dataclass, field
 from typing import Dict, List
 
 import numpy as np
+
 from luxai_s2.config import EnvConfig
 from luxai_s2.factory import Factory
 from luxai_s2.map.board import Board
 from luxai_s2.team import Team
-
 from luxai_s2.unit import Unit
-from collections import OrderedDict
-import copy
+
 
 @dataclass
 class State:
@@ -18,13 +19,12 @@ class State:
     env_steps: int
     env_cfg: EnvConfig
     board: Board = None
-    weather_schedule: np.ndarray = None
     units: Dict[str, Dict[str, Unit]] = field(default_factory=dict)
     factories: Dict[str, Dict[str, Factory]] = field(default_factory=dict)
     teams: Dict[str, Team] = field(default_factory=dict)
     global_id: int = 0
     stats: Dict[str, Dict[str, int]] = field(default_factory=dict)
-    
+
     @property
     def real_env_steps(self):
         """
@@ -47,11 +47,13 @@ def generate_unit_data(units_dict: Dict[str, Dict[str, Unit]]):
                 #     del state_dict["action_queue"]
                 units[team][unit.unit_id] = state_dict
         return units
+
     def generate_team_data(teams_dict):
         teams = dict()
         for k, v in teams_dict.items():
             teams[k] = v.state_dict()
         return teams
+
     def generate_factory_data(factories_dict):
         factories = dict()
         for team in factories_dict:
@@ -71,9 +73,9 @@ def get_obs(self):
             teams=teams,
             factories=factories,
             board=board,
-            weather_schedule=self.weather_schedule,
-            real_env_steps=self.real_env_steps
+            real_env_steps=self.real_env_steps,
         )
+
     def get_compressed_obs(self):
         # return everything on turn 0
         if self.env_steps == 0:
@@ -82,10 +84,10 @@ def get_compressed_obs(self):
         # convert lichen and lichen strains to sparse matrix format?
         del data["board"]["ore"]
         del data["board"]["ice"]
-        del data["weather_schedule"]
         if self.real_env_steps >= 0:
             del data["board"]["valid_spawns_mask"]
         return data
+
     def get_change_obs(self, prev_state):
         """
         returns sparse dicts for large matrices of where values change only
@@ -97,16 +99,20 @@ def get_change_obs(self, prev_state):
         data["board"]["lichen_strains"] = dict()
         change_indices = np.argwhere(self.board.rubble != prev_state["board"]["rubble"])
         for ind in change_indices:
-            x,y = ind[0], ind[1]
+            x, y = ind[0], ind[1]
             data["board"]["rubble"][f"{x},{y}"] = self.board.rubble[x, y]
         change_indices = np.argwhere(self.board.lichen != prev_state["board"]["lichen"])
         for ind in change_indices:
-            x,y = ind[0], ind[1]
+            x, y = ind[0], ind[1]
             data["board"]["lichen"][f"{x},{y}"] = self.board.lichen[x, y]
-        change_indices = np.argwhere(self.board.lichen_strains != prev_state["board"]["lichen_strains"])
+        change_indices = np.argwhere(
+            self.board.lichen_strains != prev_state["board"]["lichen_strains"]
+        )
         for ind in change_indices:
-            x,y = ind[0], ind[1]
-            data["board"]["lichen_strains"][f"{x},{y}"] = self.board.lichen_strains[x, y]
+            x, y = ind[0], ind[1]
+            data["board"]["lichen_strains"][f"{x},{y}"] = self.board.lichen_strains[
+                x, y
+            ]
         return data
 
     def from_obs(obs):
diff --git a/luxai_s2/luxai_s2/team.py b/luxai_s2/luxai_s2/team.py
index 8d0dcf09..75c5f32d 100644
--- a/luxai_s2/luxai_s2/team.py
+++ b/luxai_s2/luxai_s2/team.py
@@ -1,11 +1,15 @@
 from dataclasses import dataclass
 from enum import Enum
+
 from luxai_s2.config import EnvConfig
 from luxai_s2.globals import TERM_COLORS
+
 try:
     from termcolor import colored
 except:
     pass
+
+
 @dataclass
 class FactionInfo:
     color: str = "none"
@@ -32,11 +36,12 @@ def __init__(self, team_id: int, agent: str, faction: FactionTypes = None) -> No
         self.factories_to_place = 0
         self.factory_strains = []
 
-        # whether this team gets to place factories down first or not. The bid winner has this set to True. 
+        # whether this team gets to place factories down first or not. The bid winner has this set to True.
         # If tied, player_0's team has this True
         self.place_first = False
 
         self.bid = 0
+
     def state_dict(self):
         return dict(
             team_id=self.team_id,
@@ -46,8 +51,9 @@ def state_dict(self):
             metal=self.init_metal,
             factories_to_place=self.factories_to_place,
             factory_strains=self.factory_strains,
-            place_first=self.place_first
+            place_first=self.place_first,
         )
+
     def __str__(self) -> str:
         out = f"[Player {self.team_id}]"
         if TERM_COLORS:
diff --git a/luxai_s2/luxai_s2/tools/replay.py b/luxai_s2/luxai_s2/tools/replay.py
index 5be97534..e2608180 100644
--- a/luxai_s2/luxai_s2/tools/replay.py
+++ b/luxai_s2/luxai_s2/tools/replay.py
@@ -1,12 +1,13 @@
 import os
+import os.path as osp
 from typing import List
+
 from luxai_s2.replay import decode_replay_file
 from luxai_s2.state import State
-import os.path as osp
+
+
 def replay_trajectory(replay_file: str):
     decoded = decode_replay_file(replay_file)
     if "observations" in decoded:
         print("replay file already contains observations, there is nothing to do")
         return replay_file
-    
-
diff --git a/luxai_s2/luxai_s2/unit.py b/luxai_s2/luxai_s2/unit.py
index e1ea8f0c..e91d3bf2 100644
--- a/luxai_s2/luxai_s2/unit.py
+++ b/luxai_s2/luxai_s2/unit.py
@@ -1,44 +1,46 @@
+import math
 from dataclasses import dataclass
 from enum import Enum
 from typing import List
 
 import numpy as np
+
 try:
     from termcolor import colored
 except:
     pass
 from luxai_s2.config import EnvConfig, UnitConfig
-
 from luxai_s2.globals import TERM_COLORS
 from luxai_s2.map.position import Position
 from luxai_s2.team import FactionTypes, Team
 
+
 class UnitType(Enum):
     LIGHT = "Light"
     HEAVY = "Heavy"
-    
+
+
 @dataclass
 class UnitCargo:
     ice: int = 0
     ore: int = 0
     water: int = 0
     metal: int = 0
+
     def state_dict(self):
-        return dict(
-            ice=self.ice,
-            ore=self.ore,
-            water=self.water,
-            metal=self.metal
-        )
+        return dict(ice=self.ice, ore=self.ore, water=self.water, metal=self.metal)
+
 
 class Unit:
-    def __init__(self, team: Team, unit_type: UnitType, unit_id: str, env_cfg: EnvConfig) -> None:
+    def __init__(
+        self, team: Team, unit_type: UnitType, unit_id: str, env_cfg: EnvConfig
+    ) -> None:
         self.unit_type = unit_type
         self.team_id = team.team_id
         self.team = team
         self.unit_id = unit_id
         self.pos = Position(np.zeros(2, dtype=int))
-        
+
         self.cargo = UnitCargo()
         # TODO - replace with a deque perhaps?
         self.action_queue: List = []
@@ -52,27 +54,35 @@ def __str__(self) -> str:
         if TERM_COLORS:
             return colored(out, self.team.faction.value.color)
         return out
+
     def is_heavy(self) -> bool:
         return self.unit_type == UnitType.HEAVY
+
     def next_action(self):
         """
         get next action
         """
-        if len(self.action_queue) == 0: return None
+        if len(self.action_queue) == 0:
+            return None
         action = self.action_queue[0]
         return action
+
     def repeat_action(self, action):
-        # positive repeat in sequence
-        if action.repeat > 0:
-            action.repeat -= 1
-            return
-        # remove from front of queue
-        self.action_queue.pop(0)
-        # endless repeat puts action back at end of queue
-        if action.repeat == -1:
-            self.action_queue.append(action)
+        action.n -= 1
+        if action.n <= 0:
+            # remove from front of queue
+            self.action_queue.pop(0)
+            # endless repeat puts action back at end of queue
+            if action.repeat:
+                action.n = 1
+                self.action_queue.append(action)
+
     def move_power_cost(self, rubble_at_target: int):
-        return self.unit_cfg.MOVE_COST + self.unit_cfg.RUBBLE_MOVEMENT_COST * rubble_at_target
+        return math.floor(
+            self.unit_cfg.MOVE_COST
+            + self.unit_cfg.RUBBLE_MOVEMENT_COST * rubble_at_target
+        )
+
     def state_dict(self):
         return dict(
             team_id=self.team_id,
@@ -83,8 +93,10 @@ def state_dict(self):
             cargo=self.cargo.state_dict(),
             action_queue=[a.state_dict() for a in self.action_queue],
         )
+
     def add_resource(self, resource_id, amount):
-        if amount < 0: amount = 0
+        if amount < 0:
+            amount = 0
         if resource_id == 0:
             transfer_amount = min(self.cargo_space - self.cargo.ice, amount)
             self.cargo.ice += transfer_amount
@@ -101,9 +113,11 @@ def add_resource(self, resource_id, amount):
             transfer_amount = min(self.battery_capacity - self.power, amount)
             self.power += transfer_amount
         return int(transfer_amount)
+
     def sub_resource(self, resource_id, amount):
         # subtract/transfer out as much as you min(have, request)
-        if amount < 0: amount = 0
+        if amount < 0:
+            amount = 0
         if resource_id == 0:
             transfer_amount = min(self.cargo.ice, amount)
             self.cargo.ice -= transfer_amount
diff --git a/luxai_s2/luxai_s2/utils/__init__.py b/luxai_s2/luxai_s2/utils/__init__.py
index e81f4a00..bb5932b6 100644
--- a/luxai_s2/luxai_s2/utils/__init__.py
+++ b/luxai_s2/luxai_s2/utils/__init__.py
@@ -1,2 +1,2 @@
+from .animate import animate
 from .utils import is_day
-from .animate import animate
\ No newline at end of file
diff --git a/luxai_s2/luxai_s2/utils/animate.py b/luxai_s2/luxai_s2/utils/animate.py
index c7aac8bb..37896793 100644
--- a/luxai_s2/luxai_s2/utils/animate.py
+++ b/luxai_s2/luxai_s2/utils/animate.py
@@ -1,7 +1,9 @@
 def animate(imgs, _return=True, fps=10, filename="__temp__.mp4"):
     from moviepy.editor import ImageSequenceClip
+
     clip = ImageSequenceClip(imgs, fps=fps)
     clip.write_videofile(filename, verbose=False, logger=None)
     if _return:
         from IPython.display import Video
-        return Video(filename)
\ No newline at end of file
+
+        return Video(filename)
diff --git a/luxai_s2/luxai_s2/utils/utils.py b/luxai_s2/luxai_s2/utils/utils.py
index 067177b6..65158177 100644
--- a/luxai_s2/luxai_s2/utils/utils.py
+++ b/luxai_s2/luxai_s2/utils/utils.py
@@ -1,3 +1,27 @@
+from typing import List
+
 from luxai_s2.config import EnvConfig
+from luxai_s2.unit import Unit
+
+
 def is_day(config: EnvConfig, env_step):
-    return env_step % config.CYCLE_LENGTH < config.DAY_LENGTH
\ No newline at end of file
+    return env_step % config.CYCLE_LENGTH < config.DAY_LENGTH
+
+
+def get_top_two_power_units(units: List[Unit]):
+    most_power_unit: Unit = units[0]
+    most_power = -1
+    next_most_power_unit: Unit = units[1]
+    next_most_power = -1
+    for u in units:
+        if u.power > most_power:
+            next_most_power_unit = most_power_unit
+            most_power_unit = u
+            most_power = u.power
+        elif (
+            u.power >= next_most_power
+        ):  # >= check since we want to top 2 power units which can tie
+            next_most_power_unit = u
+            next_most_power = u.power
+
+    return (most_power_unit, next_most_power_unit)
diff --git a/luxai_s2/luxai_s2/version.py b/luxai_s2/luxai_s2/version.py
new file mode 100644
index 00000000..c33f6192
--- /dev/null
+++ b/luxai_s2/luxai_s2/version.py
@@ -0,0 +1,2 @@
+import pkg_resources
+__version__ = pkg_resources.require("luxai-s2")[0].version
\ No newline at end of file
diff --git a/luxai_s2/setup.py b/luxai_s2/setup.py
index ef56a00a..cfd3916d 100644
--- a/luxai_s2/setup.py
+++ b/luxai_s2/setup.py
@@ -1,15 +1,13 @@
 import os
 
-from setuptools import setup, find_packages
-
+from setuptools import find_packages, setup
 
 def read(fname):
     return open(os.path.join(os.path.dirname(__file__), fname)).read()
 
 
 setup(
-    name="luxai_s2",
-    version="2.0.0",
+    name="luxai-s2",
     author="Lux AI Challenge",
     description="The Lux AI Challenge Season 2",
     license="MIT",
@@ -17,8 +15,19 @@ def read(fname):
     url="https://github.com/Lux-AI-Challenge/Lux-Design-S2",
     packages=find_packages(exclude="kits"),
     long_description=read("../README.md"),
-    long_description_content_type='text/markdown',
-    entry_points={'console_scripts': [
-        'luxai_s2 = luxai_runner.cli:main']},
-    install_requires=['numpy', 'pygame', 'termcolor', 'matplotlib', 'pettingzoo', 'vec_noise', 'omegaconf', 'gym==0.19', 'scipy'],
+    long_description_content_type="text/markdown",
+    entry_points={"console_scripts": ["luxai-s2 = luxai_runner.cli:main"]},
+    version="2.0.0",
+    python_requires=">=3.7",
+    install_requires=[
+        "numpy",
+        "pygame",
+        "termcolor",
+        "matplotlib",
+        "pettingzoo",
+        "vec_noise",
+        "omegaconf",
+        "gym==0.19",
+        "scipy",
+    ],
 )
diff --git a/specs.md b/specs.md
index 2160addd..f4faea31 100644
--- a/specs.md
+++ b/specs.md
@@ -18,26 +18,17 @@ The rest of the document will go through the key features of this game.
 
 ## The Map
 
-The world of Lux is represented as a 2d grid. Coordinates increase east (right) and south (down). The map is always a square and is 48 tiles long. The (0, 0) coordinate is at the top left. The map has various features including [Raw Resources](#resources) (Ice, Ore), [Refined Resources](#resources) (Water, Metal), [Robots](#robots) (Light, Heavy), [Factories](#factories), [Rubble](#movement-collisions-and-rubble), and [Lichen](#lichen). The map also includes a schedule of martian weather events impacting the environment discussed below. Code wise, the coordinate (x, y) in a map feature such as rubble is indexed by `board.rubble[x][y]` for ease of use.
+The world of Lux is represented as a 2d grid. Coordinates increase east (right) and south (down). The map is always a square and is 48 tiles long. The (0, 0) coordinate is at the top left. The map has various features including [Raw Resources](#resources) (Ice, Ore), [Refined Resources](#resources) (Water, Metal), [Robots](#robots) (Light, Heavy), [Factories](#factories), [Rubble](#movement-collisions-and-rubble), and [Lichen](#lichen). Code wise, the coordinate (x, y) in a map feature such as rubble is indexed by `board.rubble[x][y]` for ease of use.
 
 Each player will start the game by bidding on an extra factory, then placing several Factories and specifying their starting resources. See the [Starting Phase](#starting-phase) for more details.
 
-## Weather Events
-
-There are 4 kinds of weather events that occur on a predetermined schedule (given to the players at the start of the game, between 3 and 5 events). These events last 1 - 30 turns with marsquakes lasting 1 - 5 turns.
-
-* Marsquake - All [Robots](#robots) generate rubble under them every turn (1/turn light, 10/turn heavy) except those on top of [Factories](#factories).
-* Cold snap - 2x power consumption
-* Dust storm - 0.5x power gain
-* Solar flare - 2x power gain
-
 ## Day/Night Cycle
 
-The Day/Night cycle consists of a 50 turn cycle, the first 30 turns being day turns, the last 20 being night turns. During the day, solar panels replenish the power of all [Robots](#robots) but during the night robots power is not recharged. [Factories](#factories) generate power each turn regardless.
+The Day/Night cycle consists of a 50 turn cycle, the first 30 turns being day turns, the last 20 being night turns. During the day, solar panels replenish the power of all [Robots](#robots) but during the night robots power is not recharged. [Factories](#factories) generate power each turn regardless with a base amount of 50 power from their nuclear reactor and additional power from connected [lichen tiles](#lichen)
 
 ## Resources
 
-There are two kinds of raw resources: Ice and Ore which can be refined by a factory into Water or Metal respectively. These resources are collected by Light or Heavy robots, then dropped off once a robot transfers them to a friendly factory, which then automatically converts them into refined resources at a constant rate. Refined resources are used for growing lichen (scoring points) as well as building more robots. Lastly, factories will process ice and ore 5 units at a time without wasting any. E.g. if a factory has 8 ore, it will refine 5 ore into 1 metal and leave 3 ore leftover; if a factory has 7 ice, it will refine 6 ice into 3 water and leave 1 ice leftover.
+There are two kinds of raw resources: Ice and Ore which can be refined by a factory into Water or Metal respectively. These resources are collected by Light or Heavy robots, then dropped off once a robot transfers them to a friendly factory, which then automatically converts them into refined resources at a constant rate. Refined resources are used for growing lichen (scoring points) as well as building more robots. Lastly, factories will process ice and ore integer units at a time without wasting any based on the processing ratio. E.g. if a factory has 8 ore, it will refine 5 ore into 1 metal and leave 3 ore leftover; if a factory has 7 ice, it will refine 4 ice into 1 water and leave 3 ice leftover.
 
 <table>
   <tr>
@@ -53,11 +44,11 @@ There are two kinds of raw resources: Ice and Ore which can be refined by a fact
   <tr>
    <td>Ice
    </td>
-   <td>50/turn
+   <td>100/turn
    </td>
    <td>Water
    </td>
-   <td>2:1
+   <td>4:1
    </td>
   </tr>
   <tr>
@@ -85,9 +76,13 @@ _Strategy Tip_: Going first is not always advantageous!
 
 [Robots](#robots) and [Factories](#factories) can perform actions each turn given certain conditions and enough power to do so. In general, all actions are simultaneously applied and are validated against the state of the game at the start of a turn. Each turn players can give an action to each factory and a queue of actions to each robot. 
 
-[Robots](#robots) always execute actions in the order of their current action queue. [Robot](#robots) actions can also be configured to be repeated `n` times or infinitely. Non-infinite action repeats mean once the action is executed `n` times succesfully, it is then removed from the action queue. When repeating infinitely `repeat = -1`, then once it is succesfully executed it is moved back to the end of the queue. Moreover, if an action in the queue fails to execute due to lack of power, that action is not removed from the queue and is kept at the front of the queue.
+[Robots](#robots) always execute actions from an action queue (limited to 20 items) while factories directly execute actions. Each robot action in the queue has an `n` value and a `repeat` value. `n` represents the number of times the robot will execute a particular action before removing it from the front of the queue and `repeat` is whether or not to replace the action to the back of the queue once it's been removed from the front.
 
-Submitting a new action queue for a robot requires the robot to use additional power to replace it's action queue. It costs an additional 1 power for Lights, an additional 10 power for Heavies. The new action queue is then stored and wipes out what was stored previously. If the robot does not have enough power, the action queue is simply not replaced.
+An action is considered executed if it's valid against current state, namely there is sufficient power to perform the action. If it's not valid, that turn it will not count towards `n`.
+
+Submitting a new action queue for a robot requires the robot to use additional power to replace it's action queue. It costs an additional 1 power for Lights, an additional 10 power for Heavies (two weight classes of robots). The new action queue is then stored and wipes out what was stored previously. If the robot does not have enough power, the action queue is simply not replaced.
+
+`repeat = True` allows you to design action queues that repeat infinitely and don't need any updates, saving power. `n` allows for generally more complex action sequences to be designed in an action queue which is limited to 20 actions.
 
 The next few sections describe the [Robots](#robots) and [Factories](#factories) in detail.
 
@@ -140,20 +135,20 @@ There are two robot types, Light and Heavy. Every robot has an action queue and
 
 ### Light and Heavy Robots
 
-Light and Heavy Robots share the same set of actions / action space. However, in general, heavy robots generally accomplish 10x more with their actions but their actions will cost 20x more power.
+Light and Heavy Robots share the same set of actions / action space. However, in general, heavy robots generally accomplish 10x more with their actions but their actions will cost more power.
 
 
 ### Actions
 
-* Move - Move the robot in one of 5 directions, North, East, South, West, Center. Moving center costs no power.
+* Move - Move the robot in one of 5 directions, North, East, South, West, Center. Moving center costs no power and is ignored.
 * Transfer -  Send any amount of a single resource-type (including power) from a robot’s cargo to an orthogonally adjacent tile or to the tile it is standing on. If a robot is on the target tile, it will receive the transferred resources up to the robot’s cargo capacity. If the target tile is a friendly factory tile, the factory receives all the transferred resources. If the receiving entity can't receive all transferred resources due to space limitations, then the overflow resources are wasted. Factories are given preference over other robots in receiving resources from transfers.
     * Algorithmically, we perform the following procedure. The environment creates transfer requests for every robot that wants to transfer and remove the specified resources from the robot’s cargo. All transfer requests are attempted to be fulfilled, and any excess caused by not enough cargo space (either robot has no space, or too many robots transferring to the same robot and go over the max capacity), is then wasted.
 * Pickup - When on top of any factory tile (there are 3x3 per factory), can pick up any amount of power or any resources. Preference is given to robots with lower robot IDs.
 * Dig - Does a number of things depending on what tile the robot is on top of
     * Rubbleless resource tile - gain raw resources (ice or ore)
     * Rubble - reduce rubble by 2 if light, 20 if heavy
-    * Lichen - reduce lichen value by 10 if light, 100 if heavy
-* Self destruct - Destroys the robot on the spot, which creates rubble.
+    * Lichen - reduce lichen value by 10 if light, 100 if heavy. If lichen value was previously > 0 and is now 0, rubble is added, 2 if light, 20 if heavy
+* Self destruct - Destroys the robot on the spot, which creates 1 rubble if light, 10 if heavy.
 * Recharge X - the robot waits until it has X power. In code, the recharge X command is not removed from the action queue until the robot has X power.
 * Repeat - This is not an explicit action but is a boolean/bit that can be added to each action. It tells the robot to append the action the robot just took to the end of the action queue. When set to False, executed actions are not appended back and are removed from the queue.
 
@@ -172,9 +167,9 @@ Action
   <tr>
    <td>Move
    </td>
-   <td>1 power
+   <td>floor(1 + 0.05 * rubble value of target square) power
    </td>
-   <td>20 power + rubble value of target square
+   <td>floor(20 + 1 * rubble value of target square) power
    </td>
   </tr>
   <tr>
@@ -198,7 +193,7 @@ Action
    </td>
    <td>5 power (2 rubble removed, 2 resources gain, 10 lichen value removed)
    </td>
-   <td>100 power (20 rubble removed, 20 resource gain, 100 lichen value removed)
+   <td>60 power (20 rubble removed, 20 resource gain, 100 lichen value removed)
    </td>
   </tr>
   <tr>
@@ -225,10 +220,13 @@ Each square on the map has a rubble value which affects how difficult that squar
 
 This environment also has robot collisions. Robots which move onto the same square on the same turn can be destroyed and add rubble to the square according the following rules:
 
-* Heavy robots that end their turn on a square with only other light robots, it will destroy all the light robots and leave the single heavy robot unaffected.
-* If two robots of the same weight end their turn on the same square, and one of them did not move in the previous turn, the stationary robot is destroyed. If both robots moved, both robots are destroyed.
+* Heavy robots that end their turn on a square with only other light robots will destroy all the light robots and leave the single heavy robot unaffected.
+* If two or more robots of the same weight end their turn on the same square, we do check the following:
+  - If only one of the units moved onto the square, the moving unit survives. 
+  - If none of the units moved onto the square, all are destroyed. 
+  - If multiple units move onto the square, then the unit with the most power survives. Moreover, that unit loses power equal to half the power of the unit with the second most power in the collision.
 
-Each light robot destroyed in this way adds 1 rubble. Each heavy robot destroyed in this way adds 10 rubble. (same values as marsquakes and self destructs).
+Each light robot destroyed in this way adds 1 rubble. Each heavy robot destroyed in this way adds 10 rubble. (same values as self destructs).
 
 Lastly, any addition of rubble onto a tile with [Lichen](#lichen) on it will automatically remove all of the lichen on that tile.
 
@@ -239,7 +237,8 @@ A factory is a building that takes up 3x3 tiles of space. Robots created from th
 Each turn a factory will automatically:
 
 * Gain 50 power (regardless of day or night)
-* Convert up to 50 ice to 25 water 
+* Gain power equal to the number of connected [lichen tiles](#lichen)
+* Convert up to 100 ice to 25 water 
 * Convert up to 50 ore to 10 metal 
 * Consume 1 water
 
@@ -249,7 +248,7 @@ Each factory can perform one of the following actions
 
 * Build a light robot
 * Build a heavy robot
-* Grow lichen - Waters lichen around the factory, costing `ceil(connected lichen tiles / 10)` water
+* Grow lichen - Waters [lichen](#lichen) around the factory, costing `ceil(connected lichen tiles / 10)` water
 
 The following is the cost to build the two classes of robots. Note that also robots when built will have their battery charged up to the power cost.
 
@@ -283,15 +282,17 @@ Robot Type
 
 ## Lichen
 
-At the end of the game, the amount of lichen on each square that a player owns is summed and whoever has a higher value wins the game. 
+At the end of the game, the amount of lichen on each square that a player owns is summed and **whoever has a higher value wins the game.**
 
-At the start, factories can perform the water action to start or continue lichen growing. Taking this action will seed lichen in all orthogonally adjacent squares to the factory if there is no rubble present (total of 3*4=12). Whenever a tile has a lichen value of 10 or more and is watered, it will spread lichen to adjacent tiles without rubble, resources, or factories and give them lichen values of 1. The amount of water consumed by the water action grows with the number of tiles with lichen on them connected to the factory according to `ceil(# connected lichen tiles / 10)`. In each tile a maximum of 100 lichen value can be stored.
+At the start, factories can perform the water action to start or continue lichen growing. Taking this action will seed lichen in all orthogonally adjacent squares to the factory if there is no rubble present (total of 3*4=12). Whenever a tile has a lichen value of 20 or more and is watered, it will spread lichen to adjacent tiles without rubble, resources, or factories and give them lichen values of 1. The amount of water consumed by the water action grows with the number of tiles with lichen on them connected to the factory according to `ceil(# connected lichen tiles / 10)`. In each tile a maximum of 100 lichen value can be stored.
 
 All factories have their own special strains of lichen that can’t mix, so lichen tiles cannot spread to tiles adjacent to lichen tiles from other factories. This is for determinism and simplified water costs.
 
-When rubble is added to a tile, that tile loses all lichen.
+Factories also **gain power** equal to the number of their own connected lichen tiles each turn.
+
+When rubble is added to a tile, that tile **loses all lichen.**
 
-If a number of lichen tiles get disconnected from your factory (due to some rubble being added to a tile or being digged out), they cannot be watered (and thus will lose 1 lichen value) until connected again through lichen tiles.
+If a number of lichen tiles get disconnected from your factory (due to some rubble being added to a tile or being dug out), they cannot be watered (and thus will lose 1 lichen value) until connected again through lichen tiles. These tiles still count to your score however.
 
 At the end of each turn, all tiles that have not been watered lose 1 lichen.
 


Test patch:
diff --git a/luxai_s2/tests/bench.py b/luxai_s2/tests/bench.py
index 5e4ac008..235c3999 100644
--- a/luxai_s2/tests/bench.py
+++ b/luxai_s2/tests/bench.py
@@ -1,9 +1,14 @@
 ### pyinstrument tests/play.py
+import copy
+
 import numpy as np
+
 from luxai_s2.env import LuxAI_S2
 from luxai_s2.replay.replay import generate_replay
-import copy
+
 spawns = None
+
+
 def policy(agent, step, obs):
     global spawns
     factory_placement_period = False
@@ -31,15 +36,17 @@ def policy(agent, step, obs):
         actions = dict()
         if step % 4 == 0 and step > 1:
             for unit_id, factory in factories.items():
-                actions[unit_id] = np.random.randint(0,2)
+                actions[unit_id] = np.random.randint(0, 2)
         else:
             for unit_id, factory in factories.items():
                 actions[unit_id] = 2
         for unit_id, unit in obs["units"][agent].items():
             # actions[unit_id] = np.array([0, np.random.randint(5), 0, 0, 0])
             # make units go to 0, 0
-            pos = unit['pos']
-            target_pos = np.array([32 + np.random.randint(-10, 10), 32 + np.random.randint(-10, 10)])
+            pos = unit["pos"]
+            target_pos = np.array(
+                [32 + np.random.randint(-10, 10), 32 + np.random.randint(-10, 10)]
+            )
             diff = target_pos - pos
             # print(pos, diff)
             direc = 0
@@ -55,26 +62,26 @@ def policy(agent, step, obs):
                     else:
                         direc = 1
             else:
-                direc = np.random.randint(0,5)
+                direc = np.random.randint(0, 5)
             actions[unit_id] = []
             for i in range(10):
                 actions[unit_id] += [np.array([0, direc, 0, 0, 0])]
         return actions
 
 
-
 if __name__ == "__main__":
     import time
+
     np.random.seed(0)
 
     env: LuxAI_S2 = LuxAI_S2(verbose=0, validate_action_space=False)
     o = env.reset(seed=0)
     render = False
-    if render: 
+    if render:
         env.render()
         time.sleep(0.1)
     states = [env.state.get_compressed_obs()]
-    
+
     s_time = time.time_ns()
     foward_pass_time = 0
     N = 10000
@@ -85,10 +92,10 @@ def policy(agent, step, obs):
         for team_id, agent in enumerate(env.possible_agents):
             all_actions[agent] = dict()
             all_actions[agent] = policy(agent, step, o)
-        
+
         o, r, d, _ = env.step(all_actions)
         step += 1
-        
+
         # states += [env.state.get_compressed_obs()]
         for agent in env.agents:
             for unit in env.state.units[agent].values():
@@ -98,8 +105,9 @@ def policy(agent, step, obs):
                 factory.cargo.water = 1000
         if np.all([d[k] for k in d]):
             o = env.reset()
-            
-            if render: env.render()
+
+            if render:
+                env.render()
             print(f"=== {i} ===")
             e_time = time.time_ns()
             print(f"FPS={step / ((e_time - s_time) * 1e-9)}")
@@ -112,6 +120,7 @@ def policy(agent, step, obs):
     # e_time = time.time_ns()
     # print(f"FPS={N / ((e_time - s_time) * 1e-9)}")
     import json
+
     def to_json(state):
         if isinstance(state, np.ndarray):
             return state.tolist()
@@ -126,6 +135,7 @@ def to_json(state):
             return out
         else:
             return state
+
     # import pickle
     # with open("replay.pkl", "wb") as f:
     #     pickle.dump(dict(states=states), f)
diff --git a/luxai_s2/tests/collision.py b/luxai_s2/tests/collision.py
index 78d0acaf..99357e4f 100644
--- a/luxai_s2/tests/collision.py
+++ b/luxai_s2/tests/collision.py
@@ -1,7 +1,9 @@
 import time
 
 import numpy as np
+
 from luxai_s2.env import LuxAI_S2
+
 if __name__ == "__main__":
     env: LuxAI_S2 = LuxAI_S2()
     o = env.reset()
@@ -13,7 +15,9 @@
     o, r, d, _ = env.step(
         {
             "player_0": dict(faction="MotherMars", spawns=np.array([[4, 4], [15, 5]])),
-            "player_1": dict(faction="AlphaStrike", spawns=np.array([[56, 55], [40, 42]])),
+            "player_1": dict(
+                faction="AlphaStrike", spawns=np.array([[56, 55], [40, 42]])
+            ),
         }
     )
     env.render()
@@ -21,6 +25,7 @@
     s_time = time.time_ns()
     N = 10000
     import ipdb
+
     for i in range(N):
         all_actions = dict()
         for team_id, agent in enumerate(env.possible_agents):
@@ -38,7 +43,7 @@
             for unit_id, unit in obs["units"][agent].items():
                 # actions[unit_id] = np.array([0, np.random.randint(5), 0, 0, 0])
                 # make units go to 0, 0
-                pos = unit['pos']
+                pos = unit["pos"]
                 target_pos = np.array([32, 32])
                 diff = target_pos - pos
                 # print(pos, diff)
@@ -70,4 +75,4 @@
         env.render()
         time.sleep(0.1)
     e_time = time.time_ns()
-    print(f"FPS={N / ((e_time - s_time) * 1e-9)}")
\ No newline at end of file
+    print(f"FPS={N / ((e_time - s_time) * 1e-9)}")
diff --git a/luxai_s2/tests/play.py b/luxai_s2/tests/play.py
index 1011990c..a95d35c8 100644
--- a/luxai_s2/tests/play.py
+++ b/luxai_s2/tests/play.py
@@ -1,23 +1,27 @@
+import copy
+
 import numpy as np
+
 from luxai_s2.env import LuxAI_S2
 from luxai_s2.replay.replay import generate_replay
-import copy
+
 if __name__ == "__main__":
     import time
+
     np.random.seed(0)
 
     env: LuxAI_S2 = LuxAI_S2(verbose=0, validate_action_space=False)
     o = env.reset()
     render = True
-    if render: 
+    if render:
         env.render()
         time.sleep(0.1)
     states = [env.state.get_compressed_obs()]
-    
+
     # u = Unit(team=Team(1, FactionTypes.MotherMars), unit_type=UnitType.HEAVY, unit_id='1s')
     # env.state.units[1].append(u)
     # observation, reward, done, info = env.last()
-    
+
     # env.render()
     # print(o, r, d)
     s_time = time.time_ns()
@@ -25,12 +29,17 @@
     N = 2000
     step = 0
     import ipdb
+
     for i in range(N):
         if step == 0:
             o, r, d, _ = env.step(
                 {
-                    "player_0": dict(faction="MotherMars", spawns=np.array([[4, 4], [15, 5]])),
-                    "player_1": dict(faction="AlphaStrike", spawns=np.array([[56, 55], [40, 42]])),
+                    "player_0": dict(
+                        faction="MotherMars", spawns=np.array([[4, 4], [15, 5]])
+                    ),
+                    "player_1": dict(
+                        faction="AlphaStrike", spawns=np.array([[56, 55], [40, 42]])
+                    ),
                 }
             )
         all_actions = dict()
@@ -45,15 +54,17 @@
             actions = dict()
             if step % 4 == 0 and step > 1:
                 for unit_id, factory in factories.items():
-                    actions[unit_id] = np.random.randint(0,2)
+                    actions[unit_id] = np.random.randint(0, 2)
             else:
                 for unit_id, factory in factories.items():
                     actions[unit_id] = 2
             for unit_id, unit in obs["units"][agent].items():
                 # actions[unit_id] = np.array([0, np.random.randint(5), 0, 0, 0])
                 # make units go to 0, 0
-                pos = unit['pos']
-                target_pos = np.array([32 + np.random.randint(-10, 10), 32 + np.random.randint(-10, 10)])
+                pos = unit["pos"]
+                target_pos = np.array(
+                    [32 + np.random.randint(-10, 10), 32 + np.random.randint(-10, 10)]
+                )
                 diff = target_pos - pos
                 # print(pos, diff)
                 direc = 0
@@ -69,15 +80,15 @@
                         else:
                             direc = 1
                 else:
-                    direc = np.random.randint(0,5)
+                    direc = np.random.randint(0, 5)
                 actions[unit_id] = []
                 for i in range(10):
                     actions[unit_id] += [np.array([0, direc, 0, 0, 0])]
             all_actions[agent] = actions
-        
+
         o, r, d, _ = env.step(all_actions)
         step += 1
-        
+
         states += [env.state.get_compressed_obs()]
         for agent in env.agents:
             for unit in env.state.units[agent].values():
@@ -87,8 +98,9 @@
                 factory.cargo.water = 1000
         if np.all([d[k] for k in d]):
             o = env.reset()
-            
-            if render: env.render()
+
+            if render:
+                env.render()
             print(f"=== {i} ===")
             e_time = time.time_ns()
             print(f"FPS={step / ((e_time - s_time) * 1e-9)}")
@@ -101,6 +113,7 @@
     # e_time = time.time_ns()
     # print(f"FPS={N / ((e_time - s_time) * 1e-9)}")
     import json
+
     def to_json(state):
         if isinstance(state, np.ndarray):
             return state.tolist()
@@ -115,6 +128,7 @@ def to_json(state):
             return out
         else:
             return state
+
     # import pickle
     # with open("replay.pkl", "wb") as f:
     #     pickle.dump(dict(states=states), f)
@@ -122,6 +136,6 @@ def to_json(state):
     # states = generate_replay(states)
     # replay = to_json(dict(states=states))
     # import ipdb;ipdb.set_trace()
-    replay=to_json(dict(states=states))
+    replay = to_json(dict(states=states))
     with open("replay.json", "w") as f:
         json.dump(replay, f)
diff --git a/luxai_s2/tests/transfer.py b/luxai_s2/tests/transfer.py
index 8d1b34b9..7e1c05dd 100644
--- a/luxai_s2/tests/transfer.py
+++ b/luxai_s2/tests/transfer.py
@@ -1,5 +1,7 @@
 import numpy as np
+
 from luxai_s2.env import LuxAI_S2
+
 if __name__ == "__main__":
     import time
 
@@ -8,6 +10,8 @@
     o, r, d, _ = env.step(
         {
             "player_0": dict(faction="MotherMars", spawns=np.array([[4, 4], [15, 5]])),
-            "player_1": dict(faction="AlphaStrike", spawns=np.array([[56, 55], [40, 42]])),
+            "player_1": dict(
+                faction="AlphaStrike", spawns=np.array([[56, 55], [40, 42]])
+            ),
         }
-    )
\ No newline at end of file
+    )