Skip to content

Commit

Permalink
Feature: update map generation
Browse files Browse the repository at this point in the history
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.

minor:
1. remove jax from dependencies. Users need to install it manually.
2. config isort to be consistent with yapf.
3. update isort to 5.12.0.
  • Loading branch information
JYuhao88 committed Jan 30, 2023
1 parent 865e9ea commit 1216d9f
Show file tree
Hide file tree
Showing 6 changed files with 217 additions and 79 deletions.
2 changes: 1 addition & 1 deletion .pre-commit-config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ repos:
- id: yapf
additional_dependencies: [toml]
- repo: https://github.com/PyCQA/isort
rev: 5.10.1
rev: 5.12.0
hooks:
- id: isort
- repo: https://github.com/PyCQA/autoflake
Expand Down
13 changes: 12 additions & 1 deletion jux/map/board.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@

from jux.config import EnvConfig, JuxBufferConfig, LuxEnvConfig
from jux.map.position import Position
from jux.map_generator.generator import GameMap, MapType, SymmetryType
from jux.map_generator.generator import GameMap, MapDistributionType, MapType, SymmetryType
from jux.utils import INT32_MAX, imax

radius = 6
Expand Down Expand Up @@ -117,12 +117,23 @@ def new(cls, seed: int, env_cfg: EnvConfig, buf_cfg: JuxBufferConfig, factories_
key=subkey,
a=jnp.array([SymmetryType.HORIZONTAL, SymmetryType.VERTICAL]),
)
key, subkey = jax.random.split(key)
map_distribution_type = jax.random.choice(
key=subkey,
a=jnp.array([
MapDistributionType.HIGH_ICE_HIGH_ORE,
MapDistributionType.HIGH_ICE_LOW_ORE,
MapDistributionType.LOW_ICE_HIGH_ORE,
MapDistributionType.LOW_ICE_LOW_ORE,
]),
)
map = GameMap.random_map(
seed=seed,
symmetry=symmetry,
map_type=map_type,
width=width,
height=height,
map_distribution_type=map_distribution_type,
)
key, subkey = jax.random.split(key)
if factories_per_team is None:
Expand Down
137 changes: 62 additions & 75 deletions jux/map_generator/generator.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
from enum import IntEnum
from functools import partial
from typing import NamedTuple, Optional, Type

Expand All @@ -12,26 +11,11 @@

from jux.config import EnvConfig, JuxBufferConfig
from jux.map_generator.flood import boundary_sum, component_sum, flood_fill
from jux.map_generator.generator_config import CaveConfig, MapDistributionType, MapType, MountainConfig
from jux.map_generator.symnoise import SymmetryNoise, SymmetryType, symmetrize
from jux.utils import INT32_MAX


class MapType(IntEnum):
CAVE = 0
CRATERS = 1
ISLAND = 2
MOUNTAIN = 3

@classmethod
def from_lux(cls, map_type: str) -> "MapType":
idx = ["Cave", "Craters", "Island", "Mountain"].index(map_type)
return cls(idx)

@classmethod
def to_lux(self) -> str:
return ["Cave", "Craters", "Island", "Mountain"][self.value]


class GameMap(NamedTuple):
rubble: jnp.int8 # int8[height, width]
ice: jnp.bool_ # bool[height, width]
Expand Down Expand Up @@ -60,14 +44,15 @@ def random_map(seed: jnp.int32 = None,
map_type: Optional[MapType] = None,
symmetry: Optional[SymmetryType] = None,
width: jnp.int32 = None,
height: jnp.int32 = None) -> "GameMap":
height: jnp.int32 = None,
map_distribution_type: jnp.int32 = None) -> "GameMap":
noise = SymmetryNoise(seed=seed, symmetry=symmetry, octaves=3)
map_rand = lax.switch(map_type, [
partial(cave, width, height),
partial(craters, width, height),
partial(island, width, height),
partial(mountain, width, height),
], symmetry, noise)
], symmetry, noise, map_distribution_type)
return GameMap.new(
map_rand.rubble,
map_rand.ice,
Expand Down Expand Up @@ -131,7 +116,14 @@ def maximum_filter(matrix, window_dimensions=(4, 4)):
return matrix


def cave(width: jnp.int32, height: jnp.int32, symmetry: SymmetryType, noise: SymmetryNoise) -> "GameMap":
def cave(
width: jnp.int32,
height: jnp.int32,
symmetry: SymmetryType,
noise: SymmetryNoise,
map_distribution_type: MapDistributionType,
) -> "GameMap":
config: CaveConfig = CaveConfig.new(map_distribution_type=map_distribution_type)
seed = noise.seed
key = jax.random.PRNGKey(seed)
key, subkey = jax.random.split(key)
Expand Down Expand Up @@ -167,19 +159,21 @@ def cave(width: jnp.int32, height: jnp.int32, symmetry: SymmetryType, noise: Sym
ice = noise.noise(x, y + 100)
ice = ice * (mask <= 1)
ice = ice * (mask != 0)
mid_mask = (ice < jnp.percentile(ice, 92)) & (ice > jnp.percentile(ice, 91))
ice = ice * (ice >= jnp.percentile(ice, 98))
ice = ice != 0
ice = mid_mask + ice * (~mid_mask)
mid_mask = (ice > jnp.percentile(ice, config.ice_mid_range[0])) & (ice < jnp.percentile(
ice, config.ice_mid_range[1]))
high_mask = (ice > jnp.percentile(ice, config.ice_high_range[0])) & (ice < jnp.percentile(
ice, config.ice_high_range[1]))
ice = mid_mask | high_mask

# Make some noisy ore, most ore is outside caves
ore = noise.noise(x, y - 100)
ore = ore * (mask != 1)
ore = ore * (mask != 0)
mid_mask = (ore < jnp.percentile(ore, 92)) & (ore > jnp.percentile(ore, 91))
ore = ore * (~(ore < jnp.percentile(ore, 98)))
ore = ore != 0
ore = mid_mask + ore * (~mid_mask)
mid_mask = (ore > jnp.percentile(ore, config.ore_mid_range[0])) & (ore < jnp.percentile(
ore, config.ore_mid_range[1]))
high_mask = (ore > jnp.percentile(ore, config.ore_high_range[0])) & (ore < jnp.percentile(
ore, config.ore_high_range[1]))
ore = mid_mask | high_mask

return GameMap(
rubble=rubble,
Expand All @@ -189,7 +183,13 @@ def cave(width: jnp.int32, height: jnp.int32, symmetry: SymmetryType, noise: Sym
)


def craters(width: jnp.int32, height: jnp.int32, symmetry: SymmetryType, noise: SymmetryNoise) -> "GameMap":
def craters(
width: jnp.int32,
height: jnp.int32,
symmetry: SymmetryType,
noise: SymmetryNoise,
map_distribution_type: MapDistributionType,
) -> "GameMap":
min_craters = jnp.maximum(2, width * height // 1000)
max_craters = jnp.maximum(4, width * height // 500)
seed = noise.seed
Expand Down Expand Up @@ -290,9 +290,8 @@ def solve_poisson(f):
cy = jnp.cos(jnp.pi * jnp.arange(ny) / (ny - 1))
f = cx[:, None] + cy[None, :] - 2

dct = jnp.divide(dct, f) * (f != 0)
dct = jnp.nan_to_num(dct, copy=True)
dct = dct * (f != 0)
dct = jnp.divide(dct, f)
dct = jnp.nan_to_num(dct, copy=True, nan=0.0, posinf=0.0, neginf=0.0)

# Return to normal space
potential = idctn1(dct)
Expand All @@ -317,7 +316,14 @@ def convolve2d_reflect(matrix, kernel):
return matrix


def mountain(width: jnp.int32, height: jnp.int32, symmetry: SymmetryType, noise: SymmetryNoise) -> "GameMap":
def mountain(
width: jnp.int32,
height: jnp.int32,
symmetry: SymmetryType,
noise: SymmetryNoise,
map_distribution_type: MapDistributionType,
) -> "GameMap":
config: MountainConfig = MountainConfig.new(map_distribution_type=map_distribution_type)
seed = noise.seed
f = jnp.zeros((height, width))

Expand Down Expand Up @@ -395,20 +401,22 @@ def mountain(width: jnp.int32, height: jnp.int32, symmetry: SymmetryType, noise:
mask = mask / jnp.amax(mask)

rubble = (100 * mask).round()
ice = (100 * mask).round()
mid_mask = (ice < jnp.percentile(ice, 50)) & (ice > jnp.percentile(ice, 48)) | (ice < jnp.percentile(
ice, 20)) & (ice > jnp.percentile(ice, 0))
ice = ice * (ice >= jnp.percentile(ice, 98))
ice = ice != 0
ice = ice * (~mid_mask) + mid_mask
ice = ice.astype(jnp.bool_)

ore = (100 * mask).round()
mid_mask = (ore < jnp.percentile(ore, 62)) & (ore > jnp.percentile(ore, 60))
ore = (ore >= jnp.percentile(ore, 83.5)) & (ore <= jnp.percentile(ore, 84.5))
ore = ore != 0
ore = ore * (~mid_mask) + mid_mask
ore = ore.astype(jnp.bool_)
ice = (100 * mask)
high_mask = (ice > jnp.percentile(ice, config.ice_high_range[0])) & (ice < jnp.percentile(
ice, config.ice_high_range[1]))
mid_mask = (ice > jnp.percentile(ice, config.ice_mid_range[0])) & (ice < jnp.percentile(
ice, config.ice_mid_range[1]))
low_mask = (ice > jnp.percentile(ice, config.ice_low_range[0])) & (ice < jnp.percentile(
ice, config.ice_low_range[1]))
ice = low_mask | mid_mask | high_mask
ore = (100 * mask)

mid_mask = (ore > jnp.percentile(ore, config.ore_mid_range[0])) & (ore < jnp.percentile(
ore, config.ore_mid_range[1]))
low_mask = (ore > jnp.percentile(ore, config.ore_low_range[0])) & (ore < jnp.percentile(
ore, config.ore_low_range[1]))

ore = low_mask | mid_mask

return GameMap(
rubble=rubble,
Expand All @@ -434,7 +442,13 @@ def convolve2d_fill(matrix: Array, kernel: Array, fillvalue: jnp.float32):
return matrix_full


def island(width: jnp.int32, height: jnp.int32, symmetry: SymmetryType, noise: SymmetryNoise) -> "GameMap":
def island(
width: jnp.int32,
height: jnp.int32,
symmetry: SymmetryType,
noise: SymmetryNoise,
map_distribution_type: MapDistributionType,
) -> "GameMap":
# at the end, 0 = island, 1 = sea (of rubble)
seed = noise.seed
key = jax.random.PRNGKey(seed)
Expand Down Expand Up @@ -496,30 +510,3 @@ def _body_func(val):
ore=ore,
symmetry=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("-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', '/', '\\')")

args = vars(parser.parse_args())
map_type = args.get("map_type", None)
symmetry = args.get("symmetry", None)
if args.get("size", None):
width = height = int(args["size"])
else:
width = height = None
if args.get("seed", None):
seed = int(args["seed"])
else:
seed = None

map_type = MapType.from_lux(map_type)
symmetry = MapType.from_lux(symmetry)
game_map = GameMap.random_map(seed=seed, symmetry=symmetry, map_type=map_type, width=width, height=height)
lux_viz(game_map)
input()
120 changes: 120 additions & 0 deletions jux/map_generator/generator_config.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,120 @@
from enum import IntEnum
from typing import NamedTuple, Tuple

from jax import lax


class MapType(IntEnum):
CAVE = 0
CRATERS = 1
ISLAND = 2
MOUNTAIN = 3

@classmethod
def from_lux(cls, map_type: str) -> "MapType":
idx = ["Cave", "Craters", "Island", "Mountain"].index(map_type)
return cls(idx)

def to_lux(self) -> str:
return ["Cave", "Craters", "Island", "Mountain"][self.value]


class MapDistributionType(IntEnum):
HIGH_ICE_HIGH_ORE = 0
HIGH_ICE_LOW_ORE = 1
LOW_ICE_HIGH_ORE = 2
LOW_ICE_LOW_ORE = 3

@classmethod
def from_lux(cls, map_distribution_type: str) -> "MapDistributionType":
idx = [
"high_ice_high_ore",
"high_ice_low_ore",
"low_ice_high_ore",
"low_ice_low_ore",
].index(map_distribution_type)
return cls(idx)

def to_lux(self) -> str:
return [
"high_ice_high_ore",
"high_ice_low_ore",
"low_ice_high_ore",
"low_ice_low_ore",
][self.value]


class CaveConfig(NamedTuple):
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.)
ore_mid_range: Tuple[float, float] = (81., 81.5)

@staticmethod
def new(map_distribution_type: MapDistributionType):
cave_config = lax.switch(
map_distribution_type,
[
lambda: CaveConfig(), # high_ice_high_ore
lambda: CaveConfig(
ore_high_range=(99.5, 100.),
ore_mid_range=(81., 81.4),
), # high_ice_low_ore
lambda: CaveConfig(
ice_high_range=(99.6, 100.),
ice_mid_range=(91., 91.5),
), # low_ice_high_ore
lambda: CaveConfig(
ore_high_range=(99.5, 100.),
ore_mid_range=(81., 81.4),
ice_high_range=(99.6, 100.),
ice_mid_range=(91., 91.5),
) # low_ice_low_ore
],
)
return cave_config


class MountainConfig(NamedTuple):

# 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.)

@staticmethod
def new(map_distribution_type: MapDistributionType):
cave_config = lax.switch(
map_distribution_type,
[
lambda: MountainConfig(), # high_ice_high_ore
lambda: MountainConfig(
ore_low_range=(61.6, 62.),
ore_mid_range=(84.5, 85.),
), # high_ice_low_ore
lambda: MountainConfig(
ice_high_range=(99.4, 100.),
ice_mid_range=(52.7, 53.),
ice_low_range=(0., 20.),
), # low_ice_high_ore
lambda: MountainConfig(
ice_high_range=(99.4, 100.),
ice_mid_range=(52.7, 53.),
ice_low_range=(0., 20.),
ore_low_range=(61.6, 62.),
ore_mid_range=(84.5, 85.),
) # low_ice_low_ore
],
)
return cave_config
Loading

0 comments on commit 1216d9f

Please sign in to comment.