parabellum 0.3.4__tar.gz → 0.5.13__tar.gz

{parabellum-0.3.4 → parabellum-0.5.13}/PKG-INFO

@@ -1,39 +1,41 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: parabellum
-Version: 0.3.4
+Version: 0.5.13
 Summary: Parabellum environment for parallel warfare simulation
-Home-page: https://github.com/syrkis/parabellum
-License: MIT
-Keywords: warfare,simulation,parallel,environment
 Author: Noah Syrkis
 Author-email: desk@syrkis.com
 Requires-Python: >=3.11,<3.12
-Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.11
+Requires-Dist: brax (>=0.12.1,<0.13.0)
+Requires-Dist: cachier (>=3.1.2,<4.0.0)
 Requires-Dist: cartopy (>=0.23.0,<0.24.0)
 Requires-Dist: contextily (>=1.6.0,<2.0.0)
-Requires-Dist: darkdetect (>=0.8.0,<0.9.0)
+Requires-Dist: distrax (>=0.1.5,<0.2.0)
 Requires-Dist: einops (>=0.8.0,<0.9.0)
+Requires-Dist: equinox (>=0.11.11,<0.12.0)
+Requires-Dist: evosax (>=0.1.6,<0.2.0)
+Requires-Dist: flashbax (>=0.1.2,<0.2.0)
+Requires-Dist: flax (>=0.10.4,<0.11.0)
 Requires-Dist: folium (>=0.17.0,<0.18.0)
 Requires-Dist: geopy (>=2.4.1,<3.0.0)
+Requires-Dist: gymnax (>=0.0.8,<0.0.9)
 Requires-Dist: ipykernel (>=6.29.5,<7.0.0)
-Requires-Dist: jax (==0.4.17)
-Requires-Dist: jaxmarl (==0.0.3)
+Requires-Dist: jax (>=0.5.0,<0.6.0)
+Requires-Dist: jax-tqdm (>=0.3.1,<0.4.0)
 Requires-Dist: jaxtyping (>=0.2.33,<0.3.0)
 Requires-Dist: jupyterlab (>=4.2.2,<5.0.0)
-Requires-Dist: moviepy (>=1.0.3,<2.0.0)
-Requires-Dist: numpy (<2)
-Requires-Dist: opencv-python (>=4.10.0.84,<5.0.0.0)
+Requires-Dist: navix (>=0.7.0,<0.8.0)
+Requires-Dist: numpy (>=2.2.3,<3.0.0)
+Requires-Dist: omegaconf (>=2.3.0,<3.0.0)
+Requires-Dist: optax (>=0.2.4,<0.3.0)
 Requires-Dist: osmnx (==2.0.0b0)
 Requires-Dist: pandas (>=2.2.2,<3.0.0)
 Requires-Dist: poetry (>=1.8.3,<2.0.0)
-Requires-Dist: pygame (>=2.5.2,<3.0.0)
 Requires-Dist: rasterio (>=1.3.10,<2.0.0)
-Requires-Dist: seaborn (>=0.13.2,<0.14.0)
 Requires-Dist: stadiamaps (>=3.2.1,<4.0.0)
 Requires-Dist: tqdm (>=4.66.4,<5.0.0)
-Project-URL: Repository, https://github.com/syrkis/parabellum
+Requires-Dist: wandb (>=0.19.7,<0.20.0)
 Description-Content-Type: text/markdown
 
 # Parabellum

parabellum-0.5.13/parabellum/__init__.py

@@ -0,0 +1,8 @@
+from . import aid, env, geo, types
+
+__all__ = [
+    "geo",
+    "env",
+    "aid",
+    "types",
+]

{parabellum-0.3.4 → parabellum-0.5.13}/parabellum/aid.py

@@ -3,10 +3,9 @@
 # by: Noah Syrkis
 
 # imports
-import os
 from collections import namedtuple
-from typing import Tuple
 import cartopy.crs as ccrs
+import jax.numpy as jnp
 
 # types
 BBox = namedtuple("BBox", ["north", "south", "east", "west"])  # type: ignore
@@ -23,5 +22,20 @@ def to_mercator(bbox: BBox) -> BBox:
 def to_platecarree(bbox: BBox) -> BBox:
     proj = ccrs.PlateCarree()
     west, south = proj.transform_point(bbox.west, bbox.south, ccrs.Mercator())
+
     east, north = proj.transform_point(bbox.east, bbox.north, ccrs.Mercator())
     return BBox(north=north, south=south, east=east, west=west)
+
+
+def obstacle_mask_fn(limit):
+    def aux(i, j):
+        xs = jnp.linspace(0, i + 1, i + j + 1)
+        ys = jnp.linspace(0, j + 1, i + j + 1)
+        cc = jnp.stack((xs, ys)).astype(jnp.int8)
+        mask = jnp.zeros((limit, limit)).at[*cc].set(1)
+        return mask
+
+    x = jnp.repeat(jnp.arange(limit), limit)
+    y = jnp.tile(jnp.arange(limit), limit)
+    mask = jnp.stack([aux(*c) for c in jnp.stack((x, y)).T])
+    return mask.astype(jnp.int8).reshape(limit, limit, limit, limit)

parabellum-0.5.13/parabellum/env.py

@@ -0,0 +1,113 @@
+# env.py
+#   parabellum env
+# by: Noah Syrkis
+
+# % Imports
+import jax.numpy as jnp
+from jax import random, Array, lax, vmap, debug
+import jax.numpy.linalg as la
+from typing import Tuple
+from functools import partial
+
+from parabellum.geo import geography_fn
+from parabellum.types import Action, State, Obs, Scene
+from parabellum import aid
+import equinox as eqx
+
+
+# %% Dataclass ################################################################
+class Env:
+    def __init__(self, cfg):
+        self.cfg = cfg
+
+    def reset(self, rng: Array, scene: Scene) -> Tuple[Obs, State]:
+        return init_fn(rng, self, scene)
+
+    def step(self, rng: Array, scene: Scene, state: State, action: Action) -> Tuple[Obs, State]:
+        return obs_fn(self, scene, state), step_fn(rng, self, scene, state, action)
+
+    @property
+    def num_units(self):
+        return sum(self.cfg.counts.allies.values()) + sum(self.cfg.counts.enemies.values())
+
+    @property
+    def num_allies(self):
+        return sum(self.cfg.counts.allies.values())
+
+    @property
+    def num_enemies(self):
+        return sum(self.cfg.counts.enemies.values())
+
+
+# %% Functions ################################################################
+@eqx.filter_jit
+def init_fn(rng: Array, env: Env, scene: Scene) -> Tuple[Obs, State]:  # initialize -----
+    keys = random.split(rng)
+    health = jnp.ones(env.num_units) * scene.unit_type_health[scene.unit_types]
+    pos = random.normal(keys[1], (scene.unit_types.size, 2)) * 2 + env.cfg.size / 2
+    state = State(unit_position=pos, unit_health=health, unit_cooldown=jnp.zeros(env.num_units))  # state --
+    return obs_fn(env, scene, state), state  # return observation and state of agents --
+
+
+@eqx.filter_jit  # knn from env.cfg never changes, so we can jit it
+def obs_fn(env, scene: Scene, state: State) -> Obs:  # return info about neighbors ---
+    distances = la.norm(state.unit_position[:, None] - state.unit_position, axis=-1)  # all dist --
+    dists, idxs = lax.approx_min_k(distances, k=env.cfg.knn)
+    mask = mask_fn(scene, state, dists, idxs)
+    health = state.unit_health[idxs] * mask
+    cooldown = state.unit_cooldown[idxs] * mask
+    unit_pos = (state.unit_position[:, None, ...] - state.unit_position[idxs]) * mask[..., None]
+    return Obs(unit_id=idxs, unit_pos=unit_pos, unit_health=health, unit_cooldown=cooldown)
+
+
+@eqx.filter_jit
+def step_fn(rng, env: Env, scene: Scene, state: State, action: Action) -> State:  # update agents ---
+    newpos = state.unit_position + action.coord * (1 - action.kinds[..., None])
+    bounds = ((newpos < 0).any(axis=-1) | (newpos >= env.cfg.size).any(axis=-1))[..., None]
+    builds = (scene.terrain.building[*newpos.astype(jnp.int32).T] > 0)[..., None]
+    newpos = jnp.where(bounds | builds, state.unit_position, newpos)  # use old pos if new is not valid
+    new_hp = blast_fn(rng, env, scene, state, action)
+    return State(unit_position=newpos, unit_health=new_hp, unit_cooldown=state.unit_cooldown)  # return -
+
+
+def blast_fn(rng, env: Env, scene: Scene, state: State, action: Action):  # update agents ---
+    dist = la.norm(state.unit_position[None, ...] - (state.unit_position + action.coord)[:, None, ...], axis=-1)
+    hits = dist <= scene.unit_type_reach[scene.unit_types][None, ...] * action.kinds[..., None]  # mask non attack act
+    damage = (hits * scene.unit_type_damage[scene.unit_types][None, ...]).sum(axis=-1)
+    return state.unit_health - damage
+
+
+# @eqx.filter_jit
+def scene_fn(cfg):  # init's a scene
+    aux = lambda key: jnp.array([x[key] for x in sorted(cfg.types, key=lambda x: x.name)])  # noqa
+    attrs = ["health", "damage", "reload", "reach", "sight", "speed"]
+    kwargs = {f"unit_type_{a}": aux(a) for a in attrs} | {"terrain": geography_fn(cfg.place, cfg.size)}
+    num_allies, num_enemies = sum(cfg.counts.allies.values()), sum(cfg.counts.enemies.values())
+    unit_teams = jnp.concat((jnp.zeros(num_allies), jnp.ones(num_enemies))).astype(jnp.int32)
+    aux = lambda t: jnp.concat([jnp.zeros(x) + i for i, x in enumerate([x[1] for x in sorted(cfg.counts[t].items())])])  # noqa
+    unit_types = jnp.concat((aux("allies"), aux("enemies"))).astype(jnp.int32)
+    mask = aid.obstacle_mask_fn(max([x["sight"] for x in cfg.types]))
+    return Scene(unit_teams=unit_teams, unit_types=unit_types, mask=mask, **kwargs)  # type: ignore
+
+
+@eqx.filter_jit
+def mask_fn(scene, state, dists, idxs):
+    mask = dists < scene.unit_type_sight[scene.unit_types][..., None]  # mask for removing hidden
+    mask = mask | obstacle_fn(scene, state.unit_position[idxs].astype(jnp.int8))
+    return mask
+
+
+@partial(vmap, in_axes=(None, 0))  # 5 x 2 # not the best name for a fn
+def obstacle_fn(scene, pos):
+    slice = slice_fn(scene, pos[0], pos)
+    return slice
+
+
+@partial(vmap, in_axes=(None, None, 0))
+def slice_fn(scene, source, target):  # returns a 10 x 10 view with unit at top left corner, and terrain downwards
+    delta = ((source - target) >= 0) * 2 - 1
+    coord = jnp.sort(jnp.stack((source, source + delta * 10)), axis=0)[0]
+    slice = lax.dynamic_slice(scene.terrain.building, coord, (scene.mask.shape[-1], scene.mask.shape[-1]))
+    slice = lax.cond(delta[0] == 1, lambda: jnp.flip(slice), lambda: slice)
+    slice = lax.cond(delta[1] == 1, lambda: jnp.flip(slice, axis=1), lambda: slice)
+    return (scene.mask[*jnp.abs(source - target)] * slice).sum() == 0

{parabellum-0.3.4 → parabellum-0.5.13}/parabellum/geo.py

@@ -3,25 +3,23 @@
 # by: Noah Syrkis
 
 # %% Imports
-from parabellum import tps
-import rasterio
 from rasterio import features, transform
+from jax import tree
 from geopy.geocoders import Nominatim
 from geopy.distance import distance
 import contextily as cx
 import jax.numpy as jnp
 import cartopy.crs as ccrs
 from jaxtyping import Array
-import numpy as np
 from shapely import box
 import osmnx as ox
 import geopandas as gpd
 from collections import namedtuple
 from typing import Tuple
 import matplotlib.pyplot as plt
-import seaborn as sns
-import os
+from cachier import cachier
 from jax.scipy.signal import convolve
+from parabellum.types import Terrain
 
 # %% Types
 Coords = Tuple[float, float]
@@ -79,27 +77,23 @@ def basemap_fn(bbox: BBox, gdf) -> Array:
     return image
 
 
-def geography_fn(place, buffer=400):
+@cachier()
+def geography_fn(place, buffer) -> Terrain:
     bbox = get_bbox(place, buffer)
     map_data = ox.features_from_bbox(bbox=bbox, tags=tags)
     gdf = gpd.GeoDataFrame(map_data)
-    gdf = gdf.clip(box(bbox.west, bbox.south, bbox.east, bbox.north)).to_crs(
-        "EPSG:3857"
-    )
+    gdf = gdf.clip(box(bbox.west, bbox.south, bbox.east, bbox.north)).to_crs("EPSG:3857")
     raster = raster_fn(gdf, shape=(buffer, buffer))
     basemap = jnp.rot90(basemap_fn(bbox, gdf), 3)
-    # 0: building", 1: "water", 2: "highway", 3: "forest", 4: "garden"
     kernel = jnp.array([[1, 1, 1], [1, 1, 1], [1, 1, 1]])
-    trans = lambda x: jnp.rot90(x, 3)
-    # <<<<<<< HEAD
-    terrain = tps.Terrain(
+    trans = lambda x: jnp.rot90(x, 3)  # noqa
+    terrain = Terrain(
         building=trans(raster[0]),
-        water=trans(
-            raster[1] - convolve(raster[1] * raster[2], kernel, mode="same") > 0
-        ),
+        water=trans(raster[1] - convolve(raster[1] * raster[2], kernel, mode="same") > 0),
         forest=trans(jnp.logical_or(raster[3], raster[4])),
         basemap=basemap,
     )
+    terrain = tree.map(lambda x: x.astype(jnp.int16), terrain)
     return terrain
 
 
@@ -128,25 +122,25 @@ def feature_fn(t, feature, gdf, shape):
 
 
 # %%
-def normalize(x):
-    return (np.array(x) - m) / (M - m)
-
-
-def get_bridges(gdf):
-    xmin, ymin, xmax, ymax = gdf.total_bounds
-    m = np.array([xmin, ymin])
-    M = np.array([xmax, ymax])
-
-    bridges = {}
-    for idx, bridge in gdf[gdf["bridge"] == "yes"].iterrows():
-        if type(bridge["name"]) == str:
-            bridges[idx[1]] = {
-                "name": bridge["name"],
-                "coords": normalize(
-                    [bridge.geometry.centroid.x, bridge.geometry.centroid.y]
-                ),
-            }
-    return bridges
+# def normalize(x):
+#     return (np.array(x) - m) / (M - m)
+
+
+# def get_bridges(gdf):
+#     xmin, ymin, xmax, ymax = gdf.total_bounds
+#     m = np.array([xmin, ymin])
+#     M = np.array([xmax, ymax])
+
+#     bridges = {}
+#     for idx, bridge in gdf[gdf["bridge"] == "yes"].iterrows():
+#         if type(bridge["name"]) == str:
+#             bridges[idx[1]] = {
+#                 "name": bridge["name"],
+#                 "coords": normalize(
+#                     [bridge.geometry.centroid.x, bridge.geometry.centroid.y]
+#                 ),
+#             }
+#     return bridges
 
 
 """

parabellum-0.5.13/parabellum/model.py

@@ -0,0 +1,6 @@
+# model.py
+#   jax model for mapping from observation to action
+# by: Noah Syrkis
+
+# Imports
+import jax.numpy as jnp

parabellum-0.5.13/parabellum/ppo.py

@@ -0,0 +1 @@
+

parabellum-0.5.13/parabellum/terrain_db.py

@@ -0,0 +1,166 @@
+# %%
+import numpy as np
+import jax.numpy as jnp
+from parabellum.types import Terrain
+
+
+# %%
+def map_raster_from_line(raster, line, size):
+    x0, y0, dx, dy = line
+    x0 = int(x0 * size)
+    y0 = int(y0 * size)
+    dx = int(dx * size)
+    dy = int(dy * size)
+    max_T = int(2**0.5 * size)
+    for t in range(max_T + 1):
+        alpha = t / float(max_T)
+        x = x0 if dx == 0 else int((1 - alpha) * x0 + alpha * (x0 + dx))
+        y = y0 if dy == 0 else int((1 - alpha) * y0 + alpha * (y0 + dy))
+        if 0 <= x < size and 0 <= y < size:
+            raster[x, y] = 1
+    return raster
+
+
+# %%
+def map_raster_from_rect(raster, rect, size):
+    x0, y0, dx, dy = rect
+    x0 = int(x0 * size)
+    y0 = int(y0 * size)
+    dx = int(dx * size)
+    dy = int(dy * size)
+    raster[x0 : x0 + dx, y0 : y0 + dy] = 1
+    return raster
+
+
+# %%
+building_color = jnp.array([201, 199, 198, 255])
+water_color = jnp.array([193, 237, 254, 255])
+forest_color = jnp.array([197, 214, 185, 255])
+empty_color = jnp.array([255, 255, 255, 255])
+
+
+def make_terrain(terrain_args, size):
+    args = {}
+    for key, config in terrain_args.items():
+        raster = np.zeros((size, size))
+        if config is not None:
+            for elem in config:
+                if "line" in elem:
+                    raster = map_raster_from_line(raster, elem["line"], size)
+                elif "rect" in elem:
+                    raster = map_raster_from_rect(raster, elem["rect"], size)
+        args[key] = jnp.array(raster.T)
+    basemap = jnp.where(
+        args["building"][:, :, None], jnp.tile(building_color, (size, size, 1)), jnp.tile(empty_color, (size, size, 1))
+    )
+    basemap = jnp.where(args["water"][:, :, None], jnp.tile(water_color, (size, size, 1)), basemap)
+    basemap = jnp.where(args["forest"][:, :, None], jnp.tile(forest_color, (size, size, 1)), basemap)
+    args["basemap"] = basemap
+    return Terrain(**args)
+
+
+# %%
+db = {
+    "blank": {"building": None, "water": None, "forest": None},
+    "F": {
+        "building": [
+            {"line": [0.25, 0.33, 0.5, 0]},
+            {"line": [0.75, 0.33, 0.0, 0.25]},
+            {"line": [0.50, 0.33, 0.0, 0.25]},
+        ],
+        "water": None,
+        "forest": None,
+    },
+    "stronghold": {
+        "building": [
+            {"line": [0.2, 0.275, 0.2, 0.0]},
+            {"line": [0.2, 0.275, 0.0, 0.2]},
+            {"line": [0.4, 0.275, 0.0, 0.2]},
+            {"line": [0.2, 0.475, 0.2, 0.0]},
+            {"line": [0.2, 0.525, 0.2, 0.0]},
+            {"line": [0.2, 0.525, 0.0, 0.2]},
+            {"line": [0.4, 0.525, 0.0, 0.2]},
+            {"line": [0.2, 0.725, 0.525, 0.0]},
+            {"line": [0.75, 0.25, 0.0, 0.2]},
+            {"line": [0.75, 0.55, 0.0, 0.19]},
+            {"line": [0.6, 0.25, 0.15, 0.0]},
+        ],
+        "water": None,
+        "forest": None,
+    },
+    "playground": {"building": [{"line": [0.5, 0.5, 0.5, 0.0]}], "water": None, "forest": None},
+    "playground2": {
+        "building": [],
+        "water": [{"rect": [0.0, 0.8, 0.1, 0.1]}, {"rect": [0.2, 0.8, 0.8, 0.1]}],
+        "forest": [{"rect": [0.0, 0.0, 1.0, 0.2]}],
+    },
+    "triangle": {
+        "building": [{"line": [0.33, 0.0, 0.0, 1.0]}, {"line": [0.66, 0.0, 0.0, 1.0]}],
+        "water": None,
+        "forest": None,
+    },
+    "u_shape": {
+        "building": [],
+        "water": [{"rect": [0.15, 0.2, 0.1, 0.5]}, {"rect": [0.4, 0.2, 0.1, 0.5]}, {"rect": [0.2, 0.2, 0.25, 0.1]}],
+        "forest": [],
+    },
+    "bridges": {
+        "building": [],
+        "water": [
+            {"rect": [0.475, 0.0, 0.05, 0.1]},
+            {"rect": [0.475, 0.15, 0.05, 0.575]},
+            {"rect": [0.475, 0.775, 0.05, 1.0]},
+            {"rect": [0.0, 0.475, 0.225, 0.05]},
+            {"rect": [0.275, 0.475, 0.45, 0.05]},
+            {"rect": [0.775, 0.475, 0.23, 0.05]},
+        ],
+        "forest": [
+            {"rect": [0.1, 0.625, 0.275, 0.275]},
+            {"rect": [0.725, 0.0, 0.3, 0.275]},
+        ],
+    },
+}
+
+# %% [raw]
+#     import matplotlib.pyplot as plt
+#     size = 100
+#     raster = np.zeros((size, size))
+#     rect = [0.475, 0., 0.05, 0.1]
+#     raster = map_raster_from_rect(raster, rect, size)
+#     rect = [0.475, 0.15, 0.05, 0.575]
+#     raster = map_raster_from_rect(raster, rect, size)
+#     rect = [0.475, 0.775, 0.05, 1.]
+#     raster = map_raster_from_rect(raster, rect, size)
+#
+#     rect = [0., 0.475, 0.225, 0.05]
+#     raster = map_raster_from_rect(raster, rect, size)
+#     rect = [0.275, 0.475, 0.45, 0.05]
+#     raster = map_raster_from_rect(raster, rect, size)
+#     rect = [0.775, 0.475, 0.23, 0.05]
+#     raster = map_raster_from_rect(raster, rect, size)
+#
+#     rect = [0.1, 0.625, 0.275, 0.275]
+#     raster = map_raster_from_rect(raster, rect, size)
+#     rect = [0.725, 0., 0.3, 0.275]
+#     raster = map_raster_from_rect(raster, rect, size)
+#
+#     plt.imshow(raster[::-1, :])
+
+# %% [markdown]
+# # Main
+
+# %%
+if __name__ == "__main__":
+    import matplotlib.pyplot as plt
+
+    # %%
+    terrain = make_terrain(db["bridges"], size=100)
+
+    # %%
+    plt.imshow(jnp.rot90(terrain.basemap))
+    bl = (39.5, 5)
+    tr = (44.5, 10)
+    plt.scatter(bl[0], 49 - bl[1])
+    plt.scatter(tr[0], 49 - tr[1], marker="+")
+
+# %%

parabellum-0.5.13/parabellum/types.py

@@ -0,0 +1,54 @@
+# types.py
+#   parabellum types
+# by: Noah Syrkis
+
+# imports
+from chex import dataclass
+from jaxtyping import Array, Bool, Float16
+
+
+# dataclasses
+@dataclass
+class State:
+    unit_position: Array
+    unit_health: Array
+    unit_cooldown: Array
+
+
+@dataclass
+class Obs:
+    unit_id: Array
+    unit_pos: Array
+    unit_health: Array
+    unit_cooldown: Array
+
+
+@dataclass
+class Action:
+    coord: Float16[Array, "... 2"]  # noqa
+    kinds: Bool[Array, "..."]
+
+
+@dataclass
+class Terrain:
+    building: Array
+    water: Array
+    forest: Array
+    basemap: Array
+
+
+@dataclass
+class Scene:
+    terrain: Terrain
+    mask: Array
+
+    unit_types: Array
+    unit_teams: Array
+
+    unit_type_health: Array
+    unit_type_damage: Array
+    unit_type_reload: Array
+
+    unit_type_reach: Array
+    unit_type_sight: Array
+    unit_type_speed: Array

parabellum-0.5.13/pyproject.toml

@@ -0,0 +1,53 @@
+[tool.poetry]
+name = "parabellum"
+version = "0.5.13"
+description = "Parabellum environment for parallel warfare simulation"
+authors = ["Noah Syrkis <desk@syrkis.com>"]
+readme = "README.md"
+
+[tool.poetry.dependencies]
+python = "^3.11,<3.12"
+jupyterlab = "^4.2.2"
+poetry = "^1.8.3"
+tqdm = "^4.66.4"
+geopy = "^2.4.1"
+osmnx = "2.0.0b0"
+rasterio = "^1.3.10"
+ipykernel = "^6.29.5"
+folium = "^0.17.0"
+pandas = "^2.2.2"
+contextily = "^1.6.0"
+einops = "^0.8.0"
+jaxtyping = "^0.2.33"
+cartopy = "^0.23.0"
+stadiamaps = "^3.2.1"
+cachier = "^3.1.2"
+equinox = "^0.11.11"
+jax = "^0.5.0"
+gymnax = "^0.0.8"
+evosax = "^0.1.6"
+distrax = "^0.1.5"
+optax = "^0.2.4"
+flax = "^0.10.4"
+numpy = "^2.2.3"
+brax = "^0.12.1"
+wandb = "^0.19.7"
+flashbax = "^0.1.2"
+navix = "^0.7.0"
+omegaconf = "^2.3.0"
+jax-tqdm = "^0.3.1"
+
+[tool.poetry.group.dev.dependencies]
+esch = { path = "../../esch" }
+purejaxrl = { path = "../purejaxrl" }
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.pyright]
+venvPath = "."
+venv = ".venv"
+
+[tool.ruff]
+line-length = 120

parabellum-0.3.4/parabellum/__init__.py

@@ -1,22 +0,0 @@
-from .env import Environment, Scenario, make_scenario, State
-from .vis import Visualizer, Skin
-from .gun import bullet_fn
-from . import vis
-from . import terrain_db
-from . import env
-from . import tps
-# from .run import run
-
-__all__ = [
-    "env",
-    "terrain_db",
-    "vis",
-    "tps",
-    "Environment",
-    "Scenario",
-    "make_scenario",
-    "State",
-    "Visualizer",
-    "Skin",
-    "bullet_fn",
-]