diff --git a/experiments/cf_nolearn.py b/experiments/cf_nolearn.py index 9fb3792a..0928a568 100644 --- a/experiments/cf_nolearn.py +++ b/experiments/cf_nolearn.py @@ -33,11 +33,7 @@ SavedProfile, is_cuda_ready, ) -from emevo.rl.ppo_normal import ( - NormalPPONet, - vmap_apply, - vmap_net, -) +from emevo.rl.ppo_normal import NormalPPONet, vmap_apply, vmap_net from emevo.spaces import BoxSpace PROJECT_ROOT = Path(__file__).parent.parent diff --git a/experiments/cf_simple.py b/experiments/cf_simple.py index 4ccb1802..dd9efe40 100644 --- a/experiments/cf_simple.py +++ b/experiments/cf_simple.py @@ -1,4 +1,5 @@ """Asexual reward evolution with Circle Foraging""" + import dataclasses import json from pathlib import Path @@ -56,11 +57,11 @@ class RewardExtractor: _max_norm: jax.Array = dataclasses.field(init=False) def __post_init__(self) -> None: - self._max_norm = jnp.sqrt(jnp.sum(self.act_space.high ** 2, axis=-1)) + self._max_norm = jnp.sqrt(jnp.sum(self.act_space.high**2, axis=-1)) def normalize_action(self, action: jax.Array) -> jax.Array: scaled = self.act_space.sigmoid_scale(action) - norm = jnp.sqrt(jnp.sum(scaled ** 2, axis=-1, keepdims=True)) + norm = jnp.sqrt(jnp.sum(scaled**2, axis=-1, keepdims=True)) return norm / self._max_norm def extract( diff --git a/noxfile.py b/noxfile.py index 071392ef..9518f9f6 100644 --- a/noxfile.py +++ b/noxfile.py @@ -6,7 +6,7 @@ import nox -SOURCES = ["src/emevo", "tests", "smoke-tests", "experiments"] +SOURCES = ["src/emevo", "tests", "smoke-tests", "experiments", "scripts"] def _sync(session: nox.Session, requirements: str) -> None: diff --git a/pyproject.toml b/pyproject.toml index 086a4c19..915d3905 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,5 +73,6 @@ select = ["E", "F", "B", "UP"] "__init__.py" = ["F401"] "src/emevo/reward_fn.py" = ["B023"] # For typer -"experiments/**/*.py" = ["B008", "UP006", "UP007"] -"smoke-tests/*.py" = ["B008", "UP006", "UP007"] \ No newline at end of file +"experiments/*.py" = ["B008", "UP006", "UP007"] +"smoke-tests/*.py" = ["B008", "UP006", "UP007"] +"scripts/*.py" = ["UP006", "UP035"] \ No newline at end of file diff --git a/scripts/make_web_data.py b/scripts/make_web_data.py new file mode 100644 index 00000000..debdee69 --- /dev/null +++ b/scripts/make_web_data.py @@ -0,0 +1,141 @@ +"""Asexual reward evolution with Circle Foraging""" + +import warnings +from pathlib import Path +from typing import List, Optional + +import numpy as np +import polars as pl +import typer + +from emevo.analysis.log_plotting import load_log + +PROJECT_ROOT = Path(__file__).parent.parent + + +def _make_stats_df(profile_and_rewards_path: Path) -> tuple[pl.DataFrame, pl.DataFrame]: + rdf = pl.read_parquet(profile_and_rewards_path) + ldf = load_log(profile_and_rewards_path.parent).cast({"unique_id": pl.Int64}) + nc_df = rdf.group_by("parent").agg(n_children=pl.col("unique_id").len()) + age_df = ( + ldf.group_by("unique_id").agg(lifetime=pl.col("unique_id").count()).collect() + ) + food_df = ldf.group_by("unique_id").agg(eaten=pl.col("n_got_food").sum()).collect() + df = ( + rdf.join( + nc_df, left_on="unique_id", right_on="parent", how="left", coalesce=True + ) + .with_columns(pl.col("n_children").replace(None, 0)) + .join(age_df, left_on="unique_id", right_on="unique_id") + .join(food_df, left_on="unique_id", right_on="unique_id") + ) + return df, ldf + + +def _agg_df( + path: Path, + start: int, + length: int, + ldf: pl.DataFrame, + ldf_offset: int, +) -> tuple[pl.DataFrame, pl.DataFrame]: + npzfile = np.load(path) + caxy = npzfile["circle_axy"][start : start + length] # (length, 200, 3) + cact = npzfile["circle_is_active"][start : start + length] # (length, 200) + saxy = npzfile["static_circle_axy"][start : start + length] + sact = npzfile["static_circle_is_active"][start : start + length] + cx_list, cy_list, ca_list = [], [], [] + sx_list, sy_list = [], [] + uniqueid_list, c_nsteps_list, s_nsteps_list = [], [], [] + for i in range(length): + active_slots = np.nonzero(cact[i]) + caxy_i = caxy[i][active_slots] + saxy_i = saxy[i][sact[i]] + + sx_list.append(saxy_i[:, 1]) + sy_list.append(saxy_i[:, 2]) + + ca_list.append(caxy_i[:, 0]) + cx_list.append(caxy_i[:, 1]) + cy_list.append(caxy_i[:, 2]) + df = ldf.filter(pl.col("step") == ldf_offset + start + i).sort("slots") + if len(df) != len(caxy_i): + warnings.warn( + "Number of active agents doesn't match" + + f"State: {len(saxy_i)} Log: {len(df)}" + + f"at step {ldf_offset + start + i}", + stacklevel=1, + ) + df = df.unique(subset="unique_id", keep="first") + df = df.filter(((pl.col("unique_id") == 0) & (pl.col("slots") != 0)).not_()) + uniqueid_list.append(df["unique_id"]) + # Num. steps + c_nsteps_list.append(df["step"]) + s_nsteps_list.append([ldf_offset + start + i] * len(saxy_i)) + + cx = np.concatenate(cx_list) + cy = np.concatenate(cy_list) + ca = np.concatenate(ca_list) + unique_id = pl.concat(uniqueid_list) + c_nsteps = pl.concat(c_nsteps_list) + + sx = np.concatenate(sx_list) + sy = np.concatenate(sy_list) + s_nsteps = np.concatenate(s_nsteps_list) + cxy_df = pl.DataFrame( + { + "angle": ca, + "x": cx, + "y": cy, + "unique_id": unique_id, + "nsteps": c_nsteps, + } + ) + sxy_df = pl.DataFrame( + { + "x": sx, + "y": sy, + "nsteps": s_nsteps, + } + ) + return cxy_df, sxy_df + + +def main( + profile_and_rewards_path: Path, + starting_points: List[int], + write_dir: Optional[Path] = None, + length: int = 100, +) -> None: + if write_dir is None: + write_dir = Path("saved-web-data") + + stats_df, ldf = _make_stats_df(profile_and_rewards_path) + stats_df.write_parquet(write_dir / "stats.parqut", compression="snappy") + + log_path = profile_and_rewards_path.parent.expanduser() + + for point in starting_points: + index = point // 1024000 + ldfi = ldf.filter( + (pl.col("step") >= point) & (pl.col("step") < point + length) + ).collect() # Offloading here for speedup + cxy_df, sxy_df = _agg_df( + log_path / f"state-{index + 1}.npz", + point - index * 1024000, + length, + ldfi, + index * 1024000, + ) + cxy_df.write_parquet( + write_dir / f"saved_cpos-{point}.parqut", + compression="snappy", + ) + sxy_df.write_parquet( + write_dir / f"saved_spos-{point}.parqut", + compression="snappy", + ) + + +if __name__ == "__main__": + typer.run(main) diff --git a/src/emevo/exp_utils.py b/src/emevo/exp_utils.py index 67cd0816..148f3a5b 100644 --- a/src/emevo/exp_utils.py +++ b/src/emevo/exp_utils.py @@ -134,7 +134,7 @@ class GopsConfig: path: str init_std: float init_mean: float - params: dict[str, float | dict[str, float]] + params: dict[str, float | dict[str, Any]] init_kwargs: dict[str, float] = dataclasses.field(default_factory=dict) def load_model(self) -> gops.Mutation | gops.Crossover: