From 77af89c3722b8ef6a078d88fa7fd66f80b6da67f Mon Sep 17 00:00:00 2001
From: LizardAPN <piskunov-05@bk.ru>
Date: Tue, 31 Mar 2026 15:39:40 +0300
Subject: [PATCH 1/5] Refactor batch size handling in Optuna tuning script

- Replaced dynamic batch size validation with a fixed list of choices to ensure consistency across trials.
- Updated the batch size selection logic to raise a `TrialPruned` exception if the chosen batch size is incompatible with the calculated buffer size, enhancing error handling during hyperparameter tuning.
---
 scripts/optuna_tune_mappo.py | 23 ++++++++---------------
 1 file changed, 8 insertions(+), 15 deletions(-)

diff --git a/scripts/optuna_tune_mappo.py b/scripts/optuna_tune_mappo.py
index 9e2553c..d69a972 100644
--- a/scripts/optuna_tune_mappo.py
+++ b/scripts/optuna_tune_mappo.py
@@ -25,17 +25,9 @@ def _study_path_slug(study_name: str) -> str:
     return s or "study"
 
 
-def _valid_batch_sizes(buffer: int) -> list[int]:
-    opts = [32, 64, 128, 256, 512, 1024]
-    valid = [b for b in opts if b <= buffer and buffer % b == 0]
-    if valid:
-        return valid
-    for b in range(8, buffer + 1):
-        if buffer % b == 0:
-            return [b]
-    if buffer >= 1:
-        return [1]
-    return []
+# Optuna requires the same categorical choices for a parameter name in every trial
+# (see CategoricalDistribution does not support dynamic value space).
+_BATCH_SIZE_CHOICES = [32, 64, 128, 256, 512, 1024]
 
 
 def main() -> None:
@@ -75,12 +67,13 @@ def objective(trial: optuna.Trial) -> float:
         n_envs = int(args.n_envs)
         n_steps = trial.suggest_categorical("n_steps", [128, 256, 512])
         buffer = n_envs * n_steps
-        valid_bs = _valid_batch_sizes(buffer)
-        if not valid_bs:
-            raise TrialPruned(f"no valid batch_size for buffer={buffer}")
 
         lr = trial.suggest_float("learning_rate", 1e-5, 1e-2, log=True)
-        batch_size = trial.suggest_categorical("batch_size", valid_bs)
+        batch_size = trial.suggest_categorical("batch_size", _BATCH_SIZE_CHOICES)
+        if batch_size > buffer or buffer % batch_size != 0:
+            raise TrialPruned(
+                f"batch_size={batch_size} incompatible with n_envs*n_steps={buffer}"
+            )
         n_epochs = trial.suggest_int("n_epochs", 4, 15)
         ent_coef = trial.suggest_float("ent_coef", 1e-4, 0.1, log=True)
         gamma = trial.suggest_float("gamma", 0.95, 0.999)

From 93b00a457b6463d6550e2448e6b6b113f565a719 Mon Sep 17 00:00:00 2001
From: LizardAPN <piskunov-05@bk.ru>
Date: Tue, 31 Mar 2026 19:05:57 +0300
Subject: [PATCH 2/5] Enhance gameplay options and introduce MAPPO self-play
 functionality

- Updated the game script to support new gameplay modes: human vs MAPPO, MAPPO vs MAPPO (self-play), and random village demo.
- Added a new function for self-play using the same MAPPO policy for both teams, improving testing and training capabilities.
- Enhanced command-line arguments to allow selection between opponent and self-play modes, ensuring clear usage instructions.
- Updated documentation and logging to reflect the new gameplay features and configurations.
---
 scripts/run_game.py                         | 82 +++++++++++++++++++--
 src/village_ai_war/play/__init__.py         |  6 +-
 src/village_ai_war/play/mappo_human_tick.py | 77 +++++++++++++++++++
 tests/test_mappo_baseline.py                | 37 +++++++++-
 4 files changed, 194 insertions(+), 8 deletions(-)

diff --git a/scripts/run_game.py b/scripts/run_game.py
index 9909e67..7857ad5 100644
--- a/scripts/run_game.py
+++ b/scripts/run_game.py
@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-"""Play: human vs MAPPO (2D), or passive demo with random village manager steps."""
+"""Play: human vs MAPPO (2D), MAPPO vs same MAPPO (no human), or random village demo."""
 
 from __future__ import annotations
 
@@ -17,7 +17,10 @@
 from village_ai_war.config_load import load_project_config
 from village_ai_war.env.game_env import GameEnv
 from village_ai_war.play.human_controls import collect_blue_bot_actions_for_tick
-from village_ai_war.play.mappo_human_tick import play_mappo_human_tick
+from village_ai_war.play.mappo_human_tick import (
+    play_mappo_human_tick,
+    play_mappo_self_play_tick,
+)
 
 
 def _resolve_ckpt(path_str: str | None) -> Path | None:
@@ -61,6 +64,10 @@ def _run_random_village_demo(
         )
         env = GameEnv(flat, mode="village", team=0, render_mode=None)
 
+    # 3D render needs a real GameState; it is only set after reset().
+    rng = np.random.default_rng(seed)
+    _obs, _ = env.reset(seed=seed)
+
     if env.render_mode == "human_3d":
         try:
             env.render()
@@ -82,12 +89,10 @@ def _run_random_village_demo(
                 )
                 env.close()
                 env = GameEnv(flat, mode="village", team=0, render_mode="human")
+                _obs, _ = env.reset(seed=seed)
             else:
                 raise
 
-    rng = np.random.default_rng(seed)
-    _obs, _ = env.reset(seed=seed)
-
     if env.render_mode is not None:
         logger.info(
             "Random village demo | render_mode={} | max_steps={} | close window or Ctrl+C to stop",
@@ -114,7 +119,7 @@ def main() -> None:
         module="pygame.pkgdata",
     )
     parser = argparse.ArgumentParser(
-        description="Village AI War: human vs MAPPO, or random village-manager demo."
+        description="Village AI War: human vs MAPPO, MAPPO self-play, or random village demo."
     )
     parser.add_argument(
         "--mappo-opponent",
@@ -122,6 +127,12 @@ def main() -> None:
         help="Path to MAPPO zip (e.g. checkpoints/bots_mappo/mappo_bot_final.zip). "
         "Human plays BLUE vs MAPPO on RED; no village AI (training-faithful micro).",
     )
+    parser.add_argument(
+        "--mappo-self-play",
+        default="",
+        help="Path to MAPPO zip: RED and BLUE both use this checkpoint (no human). "
+        "Opens 2D pygame if a display is available, otherwise runs headless.",
+    )
     parser.add_argument("--seed", type=int, default=0)
     parser.add_argument("--max-steps", type=int, default=500)
     parser.add_argument(
@@ -138,11 +149,20 @@ def main() -> None:
 
     flat = load_project_config(_ROOT)
     mappo_path = _resolve_ckpt(args.mappo_opponent or None)
+    self_play_path = _resolve_ckpt(args.mappo_self_play.strip() or None)
+
+    if str(args.mappo_opponent).strip() and str(args.mappo_self_play).strip():
+        logger.error("Use only one of --mappo-opponent or --mappo-self-play.")
+        sys.exit(1)
 
     if str(args.mappo_opponent).strip() and mappo_path is None:
         logger.error("MAPPO checkpoint not found: {}", args.mappo_opponent)
         sys.exit(1)
 
+    if str(args.mappo_self_play).strip() and self_play_path is None:
+        logger.error("MAPPO checkpoint not found: {}", args.mappo_self_play)
+        sys.exit(1)
+
     if mappo_path is not None:
         if args.human_3d:
             logger.warning("MAPPO human play uses 2D pygame only; ignoring --human-3d.")
@@ -207,6 +227,56 @@ def _render_cb(overlay_lines: tuple[str, ...] = ()) -> None:
         env.close()
         return
 
+    if self_play_path is not None:
+        if args.human_3d:
+            logger.warning(
+                "MAPPO self-play uses 2D pygame only when a display is available; "
+                "ignoring --human-3d."
+            )
+        try:
+            env = GameEnv(flat, mode="bot", team=0, render_mode="human")
+        except Exception as e:  # noqa: BLE001
+            logger.warning("Display unavailable ({}); headless MAPPO self-play", e)
+            env = GameEnv(flat, mode="bot", team=0, render_mode=None)
+
+        try:
+            mappo_model = _load_mappo_policy(self_play_path, flat)
+        except Exception as e:  # noqa: BLE001
+            logger.error("Failed to load MAPPO from {}: {}", self_play_path, e)
+            sys.exit(1)
+        logger.info("Loaded MAPPO for self-play from {}", self_play_path)
+
+        n_slots = int(flat["game"]["max_bots_for_role_change"])
+        env.reset(seed=args.seed)
+
+        if env.render_mode is not None:
+            logger.info(
+                "MAPPO vs MAPPO (same checkpoint) | max_steps={} | ESC/close to quit",
+                args.max_steps,
+            )
+        else:
+            logger.info("MAPPO self-play (headless) | max_steps={}", args.max_steps)
+
+        for t in range(args.max_steps):
+            _obs, _r, term, trunc, info = play_mappo_self_play_tick(
+                env,
+                mappo_model,
+                n_bot_slots=n_slots,
+                deterministic=args.deterministic,
+            )
+            if env.render_mode is not None:
+                env.render(
+                    overlay_lines=(
+                        f"tick={info.get('tick', '?')}  t={t}",
+                        "RED vs BLUE — same MAPPO weights",
+                    )
+                )
+            if term or trunc:
+                logger.info("Done at t={} info={}", t, info)
+                break
+        env.close()
+        return
+
     _run_random_village_demo(
         flat,
         seed=args.seed,
diff --git a/src/village_ai_war/play/__init__.py b/src/village_ai_war/play/__init__.py
index 4a7bd0e..e1ad138 100644
--- a/src/village_ai_war/play/__init__.py
+++ b/src/village_ai_war/play/__init__.py
@@ -1,6 +1,9 @@
 """Interactive play helpers (human vs AI, MAPPO observation packing)."""
 
-from village_ai_war.play.mappo_human_tick import play_mappo_human_tick
+from village_ai_war.play.mappo_human_tick import (
+    play_mappo_human_tick,
+    play_mappo_self_play_tick,
+)
 from village_ai_war.play.mappo_obs import (
     build_mappo_global_state,
     build_mappo_locals_matrix,
@@ -12,4 +15,5 @@
     "build_mappo_locals_matrix",
     "pack_mappo_observation_vector",
     "play_mappo_human_tick",
+    "play_mappo_self_play_tick",
 ]
diff --git a/src/village_ai_war/play/mappo_human_tick.py b/src/village_ai_war/play/mappo_human_tick.py
index 84043bb..6a1866f 100644
--- a/src/village_ai_war/play/mappo_human_tick.py
+++ b/src/village_ai_war/play/mappo_human_tick.py
@@ -83,3 +83,80 @@ def play_mappo_human_tick(
         learner_bot_action=None,
         learner_bot_actions=learner_bot_actions,
     )
+
+
+def play_mappo_self_play_tick(
+    env: GameEnv,
+    mappo_model: Any,
+    *,
+    n_bot_slots: int,
+    deterministic: bool = False,
+) -> tuple[Any, SupportsFloat, bool, bool, dict[str, Any]]:
+    """One tick: the same MAPPO policy controls RED and BLUE (no human).
+
+    Uses the same packed observation layout as training: fixed team-0 map / village
+    order in the global tail (:func:`build_mappo_global_state`), with local bot
+    slots built separately for team 0 and team 1.
+
+    Requires ``env.mode == \"bot\"`` and ``env.team == 0``.
+    """
+    if env.mode != "bot":
+        raise ValueError("play_mappo_self_play_tick requires GameEnv mode='bot'")
+    if env.team != 0:
+        raise ValueError("play_mappo_self_play_tick requires GameEnv team=0")
+    state = env.game_state
+    assert state is not None
+
+    gs = build_mappo_global_state(state, env._vil_obs)
+
+    def _acts_for_team(mappo_team: int) -> np.ndarray:
+        mat = build_mappo_locals_matrix(
+            state,
+            env,
+            mappo_team=mappo_team,
+            n_bot_slots=n_bot_slots,
+        )
+        packed = pack_mappo_observation_vector(mat, gs)
+        acts, _ = mappo_model.predict(packed, deterministic=deterministic)
+        return np.asarray(acts, dtype=np.int64).reshape(-1)
+
+    acts_red = _acts_for_team(0)
+    acts_blue = _acts_for_team(1)
+    if acts_red.shape[0] != n_bot_slots:
+        raise ValueError(
+            f"MAPPO expected {n_bot_slots} red actions, got {acts_red.shape[0]}"
+        )
+    if acts_blue.shape[0] != n_bot_slots:
+        raise ValueError(
+            f"MAPPO expected {n_bot_slots} blue actions, got {acts_blue.shape[0]}"
+        )
+
+    red_alive = sorted(
+        (b for b in state.villages[0].bots if b.is_alive),
+        key=lambda b: int(b.bot_id),
+    )
+    blue_alive = sorted(
+        (b for b in state.villages[1].bots if b.is_alive),
+        key=lambda b: int(b.bot_id),
+    )
+
+    env.snapshot_bot_positions_for_tick()
+    env.begin_mappo_tick()
+
+    controlled: list[tuple[int, int]] = []
+    for i, bot in enumerate(red_alive[:n_bot_slots]):
+        a = int(acts_red[i])
+        env.queue_bot_action(0, bot.bot_id, a)
+        controlled.append((bot.bot_id, a))
+    if controlled:
+        env._controlled_bot_id = controlled[0][0]
+
+    for i, bot in enumerate(blue_alive[:n_bot_slots]):
+        env.queue_bot_action(1, bot.bot_id, int(acts_blue[i]))
+
+    learner_bot_actions = {bid: ac for bid, ac in controlled}
+    return env._simulation_tick(
+        manager_action=None,
+        learner_bot_action=None,
+        learner_bot_actions=learner_bot_actions,
+    )
diff --git a/tests/test_mappo_baseline.py b/tests/test_mappo_baseline.py
index e19082b..65a4e28 100644
--- a/tests/test_mappo_baseline.py
+++ b/tests/test_mappo_baseline.py
@@ -26,7 +26,10 @@
     pack_mappo_obs_slots,
 )
 from village_ai_war.models.mappo_policy import MAPPOPolicy
-from village_ai_war.play.mappo_human_tick import play_mappo_human_tick
+from village_ai_war.play.mappo_human_tick import (
+    play_mappo_human_tick,
+    play_mappo_self_play_tick,
+)
 from village_ai_war.play.mappo_obs import (
     build_mappo_global_state,
     build_mappo_locals_matrix,
@@ -273,6 +276,38 @@ def test_play_mappo_human_tick_smoke(tmp_path: Path) -> None:
     ge.close()
 
 
+def test_play_mappo_self_play_tick_smoke(tmp_path: Path) -> None:
+    cfg = _tiny()
+    k = int(cfg["game"]["max_bots_for_role_change"])
+    venv = DummyVecEnv([lambda: MAPPOBotEnv(cfg)])
+    model = PPO(
+        MAPPOPolicy,
+        venv,
+        n_steps=64,
+        batch_size=32,
+        verbose=0,
+        policy_kwargs={"map_size": 12, "critic_hidden_dim": 64, "n_bot_slots": k},
+    )
+    save_p = tmp_path / "mappo_self"
+    model.save(str(save_p))
+    venv.close()
+
+    loaded = PPO.load(
+        str(save_p.with_suffix(".zip")),
+        device="cpu",
+        custom_objects={"lr_schedule": lambda _: 0.0, "clip_range": lambda _: 0.0},
+    )
+    ge = GameEnv(cfg, mode="bot", team=0, render_mode=None)
+    ge.reset(seed=42)
+    _o, _r, _t, _tr, _info = play_mappo_self_play_tick(
+        ge,
+        loaded,
+        n_bot_slots=k,
+        deterministic=True,
+    )
+    ge.close()
+
+
 def test_pack_mappo_obs_roundtrip_dims() -> None:
     n = 12
     loc = np.zeros((mappo_local_dim(),), dtype=np.float32)

From 39e261661ff069c929560a460b2a03c6f18a61ad Mon Sep 17 00:00:00 2001
From: LizardAPN <piskunov-05@bk.ru>
Date: Wed, 1 Apr 2026 19:05:19 +0300
Subject: [PATCH 3/5] Implement team reward mechanics and enhance bot reward
 calculations

- Introduced new team reward configurations in `bot_rewards.yaml`, including penalties and bonuses related to hunger and food security.
- Enhanced the `GameEnv` class to finalize team rewards based on controlled bots, incorporating mean/sum calculations and terminal outcomes.
- Added new methods in `BotRewardCalculator` for team-level reward shaping and terminal bonuses, improving the reward system for bot training.
- Updated tests in `test_reward_bot.py` to validate the new team reward functionalities and ensure accurate calculations for various scenarios.
---
 .gitignore                                    |   5 +-
 configs/rewards/bot_rewards.yaml              |  11 ++
 src/village_ai_war.egg-info/PKG-INFO          | 171 ------------------
 src/village_ai_war.egg-info/SOURCES.txt       |  57 ------
 .../dependency_links.txt                      |   1 -
 src/village_ai_war.egg-info/requires.txt      |   4 -
 src/village_ai_war.egg-info/top_level.txt     |   1 -
 src/village_ai_war/env/game_env.py            |  34 +++-
 src/village_ai_war/rewards/bot_reward.py      |  64 ++++++-
 tests/test_reward_bot.py                      |  52 +++++-
 10 files changed, 162 insertions(+), 238 deletions(-)
 delete mode 100644 src/village_ai_war.egg-info/PKG-INFO
 delete mode 100644 src/village_ai_war.egg-info/SOURCES.txt
 delete mode 100644 src/village_ai_war.egg-info/dependency_links.txt
 delete mode 100644 src/village_ai_war.egg-info/requires.txt
 delete mode 100644 src/village_ai_war.egg-info/top_level.txt

diff --git a/.gitignore b/.gitignore
index 823584e..d565289 100644
--- a/.gitignore
+++ b/.gitignore
@@ -11,4 +11,7 @@ outputs/
 *.egg-info
 dist/
 build/
-.cursor/
\ No newline at end of file
+.cursor/
+*.db
+*.log
+*.logs
\ No newline at end of file
diff --git a/configs/rewards/bot_rewards.yaml b/configs/rewards/bot_rewards.yaml
index a4e96e3..98c32e5 100644
--- a/configs/rewards/bot_rewards.yaml
+++ b/configs/rewards/bot_rewards.yaml
@@ -2,6 +2,17 @@
 rewards:
   bot:
     alpha: 0.6
+    reward_aggregate: mean
+    team:
+      hunger_damage_penalty: -0.025
+      fed_no_hunger_bonus: 0.06
+      food_security_coeff: 0.004
+      food_security_threshold: 80
+      food_delta_positive_coeff: 0.015
+    terminal:
+      win: 1000.0
+      loss: -1000.0
+      draw: 0.0
     warrior:
       damage_dealt: 0.15
       kill: 8.0
diff --git a/src/village_ai_war.egg-info/PKG-INFO b/src/village_ai_war.egg-info/PKG-INFO
deleted file mode 100644
index 357900b..0000000
--- a/src/village_ai_war.egg-info/PKG-INFO
+++ /dev/null
@@ -1,171 +0,0 @@
-Metadata-Version: 2.4
-Name: village-ai-war
-Version: 0.1.0
-Summary: 2D hierarchical RL environment (Gymnasium)
-Requires-Python: >=3.10
-Description-Content-Type: text/markdown
-Provides-Extra: dev
-Requires-Dist: pytest>=7.4; extra == "dev"
-Requires-Dist: ruff>=0.1; extra == "dev"
-
-# Village AI War — 2D Hierarchical RL Environment
-
-Custom Gymnasium environment for hierarchical multi-agent reinforcement learning. The **baseline is fully RL-driven**: low-level units are controlled by a learned policy (no movement heuristics), and both teams can be trained with **self-play** against pools of past checkpoints.
-
-## Architecture
-
-- **Bot agents (low level):** One **role-conditioned** PPO policy (`RoleConditionedPolicy`) — shared backbone plus a learned role embedding from the observation one-hot (see `BotObsBuilder`). All roles share weights.
-- **Village agent (high level):** Strategic manager — **MaskablePPO** with invalid-action masking (`MultiInputPolicy` on dict observations).
-- **Self-play:** Stage 1 samples opponents from `checkpoints/pool/bots/`; stage 2 samples village opponents from `checkpoints/pool/village/`. Empty pools fall back to random opponent actions.
-- **Unified training:** One Hydra entry point (`training=train_unified`, `training.stage=0`) alternates PPO on bots and MaskablePPO on the red manager. Each environment step matches village self-play order (all bots act, then both managers). Blue bots and blue manager are sampled from the same pools; the non-training partner on red is frozen from the last saved checkpoint until the next phase.
-- **Reward shaping:** Dynamic global reward modes controlled by the village agent (unchanged).
-
-## Quick Start
-
-```bash
-cd VillageAI_War
-python3 -m venv .venv && source .venv/bin/activate   # Windows: .venv\Scripts\activate
-pip install -r requirements.txt
-pip install -e .
-python scripts/run_game.py
-```
-
-### Watching trained policies
-
-[`scripts/run_game.py`](scripts/run_game.py) can load the same checkpoints produced by training. If a path is missing or fails to load, that component falls back to random valid actions (same as the old demo).
-
-```bash
-# Defaults: checkpoints/village/village_final.zip, checkpoints/bots/bot_final.zip
-python scripts/run_game.py
-
-# Explicit paths and deterministic manager actions
-python scripts/run_game.py \
-  --village-checkpoint checkpoints/village/village_final.zip \
-  --opponent-village-checkpoint checkpoints/pool/village/village_iter5.zip \
-  --bot-checkpoint checkpoints/bots/bot_final.zip \
-  --deterministic --seed 42 --max-steps 2000
-
-# Artifacts from unified training
-python scripts/run_game.py \
-  --village-checkpoint checkpoints/unified/village_final.zip \
-  --bot-checkpoint checkpoints/unified/bot_final.zip \
-  --deterministic --seed 42
-```
-
-When any village, opponent, or bot checkpoint is loaded (or you pass a path that exists), the viewer runs **both** village managers each tick (trained or random per side), then resolves the tick in the same order as self-play training (`run_bots_then_village_decisions` in [`GameEnv`](src/village_ai_war/env/game_env.py)). If **no** checkpoints load, behavior matches the legacy script: one random manager action per step for team 0 only.
-
-The pygame **human** window includes numeric **row/column axes**, a **legend** (terrain colors, harvest hints `w`/`s`/`f`, unit roles and team rings, building abbreviations), and a **bottom HUD** (tick, winner, resources, population, global mode per team). `rgb_array` mode is still the raw map only for headless frames.
-
-**3D view** (OpenGL via [moderngl](https://github.com/moderngl/moderngl)): run `python scripts/run_game.py --human-3d` for an extruded terrain board, **building silhouettes** (town hall roof, tower spire, farm silo, etc.), **bots as team disk + role sphere + team ring**, and a **Russian legend panel** on the right (including a **live HUD**: tick, wood/stone/food, alive bots / pop cap, village AI mode per team). **Camera:** left-drag on the map to orbit (yaw/pitch); arrow keys nudge the camera; optional idle spin via `auto_rotate_deg_per_sec` (default `0`). Tuning: `window_width_3d`, `legend_width_3d`, `camera_dist_scale`, `orbit_mouse_sensitivity`, `orbit_key_deg_per_sec`, … in [`configs/default.yaml`](configs/default.yaml). In code, use `GameEnv(..., render_mode="human_3d")` or `render_mode="rgb_array_3d"` for full-window RGB captures (map + legend).
-
-On **Linux**, Mesa packages expose `libGL.so.1` (not `libGL.so`); this project patches library loading so moderngl finds them. If you still see GL errors, install dri drivers: `sudo apt install -y libgl1-mesa-dri libegl1`. On **WSL2** you need a working GUI stack (WSLg); without it, use 2D or run from native Windows. If 3D fails, `run_game.py` falls back to the 2D pygame window and logs a warning.
-
-### Training (Hydra)
-
-The default config composes `training: train_bots_selfplay` (see [`configs/default.yaml`](configs/default.yaml)). Stages are selected with `training.stage` (`0` = unified, `1`–`3` = legacy pipeline).
-
-**Stage 1 — bot self-play (role-conditioned PPO)**
-
-```bash
-python scripts/run_training.py training.stage=1
-```
-
-Uses `game.initial_bots: 1` in [`configs/training/train_bots_selfplay.yaml`](configs/training/train_bots_selfplay.yaml) so each team has one controllable unit per episode (reproducible with `SubprocVecEnv`; multi-bot stage 1 would need extra machinery such as checkpoint sync or `DummyVecEnv`-only runs).
-
-**Stage 2 — village self-play (MaskablePPO, frozen bot policy)**
-
-Requires a stage-1 artifact at `checkpoints/bots/bot_final.zip`.
-
-```bash
-python scripts/run_training.py training=train_village_selfplay training.stage=2
-```
-
-**Stage 3 — joint fine-tuning**
-
-Loads `checkpoints/village/village_final.zip` when present; bots use `checkpoints/bots/bot_final.zip` when present (otherwise random bot actions).
-
-```bash
-python scripts/run_training.py training=train_joint training.stage=3
-```
-
-**Unified training (recommended)** — `training.stage=0`, single process alternating two learners
-
-No prerequisite checkpoints. Team 0’s bot policy (PPO) and red manager (MaskablePPO) are updated in turns; after each phase the trainer saves so the other phase loads a frozen partner from disk. Opponents use the same self-play pools as stages 1–2; empty pools still fall back to random valid actions.
-
-```bash
-python scripts/run_training.py training=train_unified
-```
-
-You can also set `training.stage=0` with another `training=` group; built-in defaults for `unified.*` still apply, but prefer `training=train_unified` so values in [`configs/training/train_unified.yaml`](configs/training/train_unified.yaml) are loaded.
-
-Configure macro steps with `unified.bot_steps_per_turn`, `unified.village_steps_per_turn`, `unified.n_cycles`, optional `unified.first_phase` (`bot` or `village`), and `unified.push_to_pool` (append snapshots to `checkpoints/pool/bots` and `.../village`) in [`configs/training/train_unified.yaml`](configs/training/train_unified.yaml). Each phase logs an estimated **SB3 iteration count** (`n_steps × n_envs` env-steps per iteration). `unified.progress_bar: true` enables Stable-Baselines’ tqdm progress bar for the current `learn()` (requires `tqdm` and `rich` in [`requirements.txt`](requirements.txt)). `unified.sb3_verbose` controls SB3’s stdout tables (`0` by default when using the progress bar). `unified.progress_log_interval_sec` throttles **Loguru** lines with progress, ETA, env-steps/s, and wall time of the previous full SB3 iteration (see `run_training.log` under the Hydra run directory). Set `unified.plot_metrics_on_finish: true` to write PNG grids of TensorBoard scalars into the same Hydra output folder when the run finishes (or run [`scripts/plot_tensorboard_scalars.py`](scripts/plot_tensorboard_scalars.py) manually). Outputs live under `checkpoints/unified/` (`bot_final.zip`, `village_final.zip`, plus per-cycle checkpoints and `bot_latest` / `village_latest` stems used between phases).
-
-The bot phase uses `DummyVecEnv` only so every sub-env shares the in-process `bot_policy_holder`; do not use `SubprocVecEnv` for that phase. Default `game.initial_bots: 1` matches stage 1; more red bots require a live model in the holder for the extra units.
-
-Stages 1–3 remain a supported alternative pipeline.
-
-**Useful overrides**
-
-```bash
-python scripts/run_training.py training.total_timesteps=2000 training.n_envs=1
-```
-
-### Metrics (TensorBoard)
-
-With `logging.use_tensorboard: true` (default) and `tensorboard` installed, stage 1 and 2 write training scalars under `logs/bots/` and `logs/village/`; unified training writes under `logs/unified_bots/` and `logs/unified_village/`. Periodic evaluation logs `eval/mean_reward` (and related fields) under `logs/bots_eval/` and `logs/village_eval/` when `training.eval_freq > 0` (default `10000` **environment timesteps** between evals; internally scaled by `n_envs` per Stable-Baselines3). The unified config sets `eval_freq: 0` by default (no separate eval pass yet). Tune with `training.n_eval_episodes`.
-
-```bash
-tensorboard --logdir logs/
-```
-
-After unified training (or anytime), generate static PNG grids from the latest run in each subfolder:
-
-```bash
-python scripts/plot_tensorboard_scalars.py --log-root logs
-```
-
-Default outputs: `logs/plots/unified_bots_scalars.png` and `logs/plots/unified_village_scalars.png` (or under the Hydra run directory when `unified.plot_metrics_on_finish` is enabled).
-
-**Best vs last checkpoint:** when evaluation is enabled and at least one eval run produced a best model, `checkpoints/bots/bot_final.zip` and `checkpoints/village/village_final.zip` are copies of the best eval checkpoint (also saved as `bot_best.zip` / `village_best.zip`). The last weights after all self-play iterations are kept as `bot_last.zip` / `village_last.zip`. Set `training.eval_freq=0` to keep the previous behavior (final = last iteration only). Stage 3 joint training does not add a separate eval pass yet; it still saves `checkpoints/joint/joint_final.zip` from the end of the run. Unified training always ends with `bot_final.zip` / `village_final.zip` as the last full save after all cycles (no best-model selection unless you add eval later).
-
-**Evaluation**
-
-```bash
-python scripts/evaluate.py
-```
-
-## Checkpoints (typical layout)
-
-| Path | Contents |
-|------|-----------|
-| `checkpoints/bots/bot_final.zip` | Stage 1 policy (best eval mean reward when `training.eval_freq > 0`, else last iteration) |
-| `checkpoints/pool/bots/*.zip` | Historical bot policies for self-play |
-| `checkpoints/village/village_final.zip` | Stage 2 manager (same best-vs-last rule as bots) |
-| `checkpoints/pool/village/*.zip` | Historical village policies for self-play |
-| `checkpoints/joint/joint_final.zip` | Stage 3 output |
-| `checkpoints/unified/bot_final.zip` | Unified loop bot policy (last save after all cycles) |
-| `checkpoints/unified/village_final.zip` | Unified loop village policy (last save after all cycles) |
-| `checkpoints/unified/bot_latest.zip` / `village_latest.zip` | Latest weights exchanged between alternating phases |
-| `checkpoints/unified/bot_cycle*.zip` / `village_cycle*.zip` | Periodic `CheckpointCallback` snapshots during unified runs |
-
-## Project layout (RL baseline)
-
-- [`src/village_ai_war/env/game_env.py`](src/village_ai_war/env/game_env.py) — `step`, `step_with_opponent`, `step_village_only`, optional `game.bot_rl_checkpoint` / `training.bot_checkpoint` for frozen bot PPO
-- [`src/village_ai_war/models/role_conditioned_policy.py`](src/village_ai_war/models/role_conditioned_policy.py)
-- [`src/village_ai_war/training/self_play_env.py`](src/village_ai_war/training/self_play_env.py) — `SelfPlayBotEnv`, `SelfPlayVillageEnv`, `UnifiedBotSelfPlayEnv`
-- [`src/village_ai_war/training/train_bots_selfplay.py`](src/village_ai_war/training/train_bots_selfplay.py), [`train_village_selfplay.py`](src/village_ai_war/training/train_village_selfplay.py), [`train_joint.py`](src/village_ai_war/training/train_joint.py), [`train_unified.py`](src/village_ai_war/training/train_unified.py), [`tensorboard_plots.py`](src/village_ai_war/training/tensorboard_plots.py), [`plot_tensorboard_scalars.py`](scripts/plot_tensorboard_scalars.py)
-
-Legacy trainers [`train_bots.py`](src/village_ai_war/training/train_bots.py) and [`train_village.py`](src/village_ai_war/training/train_village.py) are not used by [`scripts/run_training.py`](scripts/run_training.py).
-
-## Project status
-
-- [x] Data structures (Pydantic)
-- [x] Map generator
-- [x] Economy / combat / building systems
-- [x] Observation builders and action masker
-- [x] GameEnv (Gymnasium), no bot heuristics
-- [x] Role-conditioned bot policy + PPO self-play (stage 1)
-- [x] Village MaskablePPO self-play with RL bots (stage 2)
-- [x] Joint fine-tuning with RL bots (stage 3)
-- [x] Unified training loop (alternating bot PPO + village MaskablePPO)
-- [x] Pygame renderer
diff --git a/src/village_ai_war.egg-info/SOURCES.txt b/src/village_ai_war.egg-info/SOURCES.txt
deleted file mode 100644
index 3ba5e46..0000000
--- a/src/village_ai_war.egg-info/SOURCES.txt
+++ /dev/null
@@ -1,57 +0,0 @@
-README.md
-pyproject.toml
-src/village_ai_war/__init__.py
-src/village_ai_war/config_load.py
-src/village_ai_war/exceptions.py
-src/village_ai_war.egg-info/PKG-INFO
-src/village_ai_war.egg-info/SOURCES.txt
-src/village_ai_war.egg-info/dependency_links.txt
-src/village_ai_war.egg-info/requires.txt
-src/village_ai_war.egg-info/top_level.txt
-src/village_ai_war/agents/__init__.py
-src/village_ai_war/agents/action_masker.py
-src/village_ai_war/agents/bot_obs_builder.py
-src/village_ai_war/agents/village_action_space.py
-src/village_ai_war/agents/village_obs_builder.py
-src/village_ai_war/env/__init__.py
-src/village_ai_war/env/building_system.py
-src/village_ai_war/env/combat_system.py
-src/village_ai_war/env/economy_system.py
-src/village_ai_war/env/game_env.py
-src/village_ai_war/env/map_generator.py
-src/village_ai_war/models/__init__.py
-src/village_ai_war/models/role_conditioned_policy.py
-src/village_ai_war/rendering/__init__.py
-src/village_ai_war/rendering/moderngl_3d_renderer.py
-src/village_ai_war/rendering/pygame_renderer.py
-src/village_ai_war/rewards/__init__.py
-src/village_ai_war/rewards/bot_reward.py
-src/village_ai_war/rewards/global_reward.py
-src/village_ai_war/rewards/village_reward.py
-src/village_ai_war/state/__init__.py
-src/village_ai_war/state/bot_state.py
-src/village_ai_war/state/constants.py
-src/village_ai_war/state/game_state.py
-src/village_ai_war/state/village_state.py
-src/village_ai_war/training/__init__.py
-src/village_ai_war/training/pool_manager.py
-src/village_ai_war/training/progress_callback.py
-src/village_ai_war/training/self_play_env.py
-src/village_ai_war/training/tensorboard_plots.py
-src/village_ai_war/training/train_bots.py
-src/village_ai_war/training/train_bots_selfplay.py
-src/village_ai_war/training/train_joint.py
-src/village_ai_war/training/train_unified.py
-src/village_ai_war/training/train_village.py
-src/village_ai_war/training/train_village_selfplay.py
-tests/test_action_masker.py
-tests/test_combat_system.py
-tests/test_economy_system.py
-tests/test_game_state.py
-tests/test_pool_manager.py
-tests/test_render_import.py
-tests/test_reward_bot.py
-tests/test_reward_village.py
-tests/test_role_conditioned_policy.py
-tests/test_smoke_game_env.py
-tests/test_unified_training_smoke.py
\ No newline at end of file
diff --git a/src/village_ai_war.egg-info/dependency_links.txt b/src/village_ai_war.egg-info/dependency_links.txt
deleted file mode 100644
index 8b13789..0000000
--- a/src/village_ai_war.egg-info/dependency_links.txt
+++ /dev/null
@@ -1 +0,0 @@
-
diff --git a/src/village_ai_war.egg-info/requires.txt b/src/village_ai_war.egg-info/requires.txt
deleted file mode 100644
index 20600e1..0000000
--- a/src/village_ai_war.egg-info/requires.txt
+++ /dev/null
@@ -1,4 +0,0 @@
-
-[dev]
-pytest>=7.4
-ruff>=0.1
diff --git a/src/village_ai_war.egg-info/top_level.txt b/src/village_ai_war.egg-info/top_level.txt
deleted file mode 100644
index 988f68b..0000000
--- a/src/village_ai_war.egg-info/top_level.txt
+++ /dev/null
@@ -1 +0,0 @@
-village_ai_war
diff --git a/src/village_ai_war/env/game_env.py b/src/village_ai_war/env/game_env.py
index b5c3c95..965d07f 100644
--- a/src/village_ai_war/env/game_env.py
+++ b/src/village_ai_war/env/game_env.py
@@ -520,6 +520,7 @@ def _run_simulation_phase(
         if self.mode == "bot" and learner_bot_actions is not None:
             mode = state.villages[self.team].global_reward_mode
             reward = 0.0
+            n_controlled = 0
             for bid, act in learner_bot_actions.items():
                 bot = next(
                     (b for v in state.villages for b in v.bots if b.bot_id == bid),
@@ -528,6 +529,10 @@ def _run_simulation_phase(
                 if bot is not None:
                     bev = self._bot_events_for(bot, merged, int(act), state)
                     reward += float(BotRewardCalculator.compute(bev, bot, mode, self.config))
+                    n_controlled += 1
+            reward = self._finalize_bot_team_reward(
+                reward, n_controlled, merged, state, terminated, truncated
+            )
         elif self.mode == "bot" and learner_bot_action is not None:
             bot = next(
                 (
@@ -542,8 +547,13 @@ def _run_simulation_phase(
                 mode = state.villages[self.team].global_reward_mode
                 bev = self._bot_events_for(bot, merged, learner_bot_action, state)
                 reward = float(BotRewardCalculator.compute(bev, bot, mode, self.config))
+                reward = self._finalize_bot_team_reward(
+                    reward, 1, merged, state, terminated, truncated
+                )
             else:
-                reward = 0.0
+                reward = self._finalize_bot_team_reward(
+                    0.0, 0, merged, state, terminated, truncated
+                )
         else:
             reward = float(
                 VillageRewardCalculator.compute(
@@ -767,6 +777,28 @@ def _nearest_dist(pos: tuple[int, int], targets: list[tuple[int, int]]) -> float
         px, py = pos
         return float(min(abs(px - tx) + abs(py - ty) for tx, ty in targets))
 
+    def _finalize_bot_team_reward(
+        self,
+        per_bot_sum: float,
+        n_controlled: int,
+        merged: Mapping[str, Any],
+        state: GameState,
+        terminated: bool,
+        truncated: bool,
+    ) -> float:
+        """Mean/sum over controlled bots, then team shaping and terminal win/loss."""
+        rcfg = self.config["rewards"]["bot"]
+        agg = str(rcfg.get("reward_aggregate", "mean")).lower()
+        r = per_bot_sum
+        if agg == "mean" and n_controlled > 0:
+            r /= float(n_controlled)
+        vil = state.villages[self.team]
+        r += BotRewardCalculator.team_addon(merged, vil, self.config)
+        r += BotRewardCalculator.terminal_addon(
+            self.config, terminated or truncated, self.team, state.winner
+        )
+        return r
+
     def _bot_events_for(
         self,
         bot: BotState,
diff --git a/src/village_ai_war/rewards/bot_reward.py b/src/village_ai_war/rewards/bot_reward.py
index 443b4d6..fdfa08b 100644
--- a/src/village_ai_war/rewards/bot_reward.py
+++ b/src/village_ai_war/rewards/bot_reward.py
@@ -6,7 +6,7 @@
 from typing import Any
 
 from village_ai_war.rewards.global_reward import mode_coefficient
-from village_ai_war.state import BotState, GlobalRewardMode
+from village_ai_war.state import BotState, GlobalRewardMode, VillageState
 
 
 class BotRewardCalculator:
@@ -49,3 +49,65 @@ def compute(
         )
         global_part = gmult * float(events.get("global_scale", 1.0))
         return float(alpha) * local + (1.0 - float(alpha)) * global_part
+
+    @staticmethod
+    def team_addon(
+        merged: Mapping[str, Any],
+        village: VillageState,
+        config: Mapping[str, Any],
+    ) -> float:
+        """Team-level shaping once per tick (hunger, food stock, net food gain)."""
+        rcfg = config["rewards"]["bot"]
+        raw = rcfg.get("team")
+        if not isinstance(raw, Mapping):
+            return 0.0
+        team_cfg: Mapping[str, Any] = raw
+        team = int(village.team)
+        r = 0.0
+
+        hd_map = merged.get("hunger_damage")
+        hd = int(hd_map.get(team, 0)) if isinstance(hd_map, Mapping) else 0
+        hp = float(team_cfg.get("hunger_damage_penalty", 0.0))
+        if hp != 0.0 and hd > 0:
+            r += hp * float(hd)
+
+        alive = sum(1 for b in village.bots if b.is_alive)
+        fbonus = float(team_cfg.get("fed_no_hunger_bonus", 0.0))
+        if alive > 0 and hd == 0 and fbonus != 0.0:
+            r += fbonus
+
+        coeff = float(team_cfg.get("food_security_coeff", 0.0))
+        if coeff != 0.0:
+            thresh = int(team_cfg.get("food_security_threshold", 0))
+            food = int(village.resources.food)
+            if food > thresh:
+                r += coeff * float(food - thresh) / 1000.0
+
+        fd_map = merged.get("food_delta")
+        if isinstance(fd_map, Mapping):
+            fd = int(fd_map.get(team, 0))
+            fdc = float(team_cfg.get("food_delta_positive_coeff", 0.0))
+            if fdc != 0.0 and fd > 0:
+                r += fdc * float(fd) / 100.0
+
+        return r
+
+    @staticmethod
+    def terminal_addon(
+        config: Mapping[str, Any],
+        episode_done: bool,
+        team: int,
+        winner: int | None,
+    ) -> float:
+        """Sparse win/loss/draw bonus for bot training (mirrors village terminal intent)."""
+        if not episode_done:
+            return 0.0
+        rcfg = config["rewards"]["bot"]
+        tc = rcfg.get("terminal")
+        if not isinstance(tc, Mapping):
+            return 0.0
+        if winner is None:
+            return float(tc.get("draw", 0.0))
+        if int(winner) == int(team):
+            return float(tc.get("win", 0.0))
+        return float(tc.get("loss", 0.0))
diff --git a/tests/test_reward_bot.py b/tests/test_reward_bot.py
index fef6ac4..a7f8b14 100644
--- a/tests/test_reward_bot.py
+++ b/tests/test_reward_bot.py
@@ -2,8 +2,10 @@
 
 from typing import Any
 
+import pytest
+
 from village_ai_war.rewards.bot_reward import BotRewardCalculator
-from village_ai_war.state import BotState, GlobalRewardMode, Role
+from village_ai_war.state import BotState, GlobalRewardMode, ResourceStock, Role, VillageState
 
 
 def _cfg() -> dict[str, Any]:
@@ -31,6 +33,54 @@ def _cfg() -> dict[str, Any]:
     }
 
 
+def _cfg_team_terminal() -> dict[str, Any]:
+    c = _cfg()
+    c["rewards"]["bot"]["team"] = {
+        "hunger_damage_penalty": -0.1,
+        "fed_no_hunger_bonus": 0.5,
+        "food_security_coeff": 1.0,
+        "food_security_threshold": 100,
+        "food_delta_positive_coeff": 1.0,
+    }
+    c["rewards"]["bot"]["terminal"] = {"win": 10.0, "loss": -10.0, "draw": -1.0}
+    return c
+
+
+def test_team_addon_hunger_penalty() -> None:
+    vil = VillageState(
+        team=0,
+        bots=[BotState(bot_id=0, team=0, role=Role.WARRIOR, position=(0, 0))],
+    )
+    merged: dict[str, Any] = {"hunger_damage": {0: 10}, "food_delta": {0: 0}}
+    r = BotRewardCalculator.team_addon(merged, vil, _cfg_team_terminal())
+    assert r == -0.1 * 10.0
+
+
+def test_team_addon_fed_and_food_delta() -> None:
+    vil = VillageState(
+        team=0,
+        resources=ResourceStock(food=50),
+        bots=[BotState(bot_id=0, team=0, role=Role.WARRIOR, position=(0, 0))],
+    )
+    merged: dict[str, Any] = {"hunger_damage": {0: 0}, "food_delta": {0: 50}}
+    r = BotRewardCalculator.team_addon(merged, vil, _cfg_team_terminal())
+    assert r == pytest.approx(0.5 + 1.0 * 50.0 / 100.0)
+
+
+def test_team_addon_no_team_section() -> None:
+    vil = VillageState(team=0, bots=[])
+    assert BotRewardCalculator.team_addon({"hunger_damage": {0: 99}}, vil, _cfg()) == 0.0
+
+
+def test_terminal_addon() -> None:
+    cfg = _cfg_team_terminal()
+    assert BotRewardCalculator.terminal_addon(cfg, False, 0, 0) == 0.0
+    assert BotRewardCalculator.terminal_addon(cfg, True, 0, 0) == 10.0
+    assert BotRewardCalculator.terminal_addon(cfg, True, 0, 1) == -10.0
+    assert BotRewardCalculator.terminal_addon(cfg, True, 0, None) == -1.0
+    assert BotRewardCalculator.terminal_addon(_cfg(), True, 0, 0) == 0.0
+
+
 def test_warrior_damage_local() -> None:
     bot = BotState(bot_id=0, team=0, role=Role.WARRIOR, position=(0, 0))
     r = BotRewardCalculator.compute(

From 0c684edac3e164d867fddf1be45bd8208f8c1115 Mon Sep 17 00:00:00 2001
From: LizardAPN <piskunov-05@bk.ru>
Date: Wed, 1 Apr 2026 21:32:22 +0300
Subject: [PATCH 4/5] Enhance MAPPO training configurations and metrics
 tracking

- Added `max_grad_norm` parameter to training configuration for improved gradient clipping.
- Updated Optuna tuning script to include new choices for `n_steps`, `gamma/gae` presets, and `checkpoint_interval`, enhancing hyperparameter search capabilities.
- Introduced `win_townhall_frac` metric in the episode metrics callback to track the fraction of wins via opponent town hall destruction, improving performance analysis.
- Modified training function to return `win_townhall_frac` alongside existing metrics, providing more insights into training outcomes.
---
 configs/training/train_mappo_bots.yaml        |  1 +
 scripts/optuna_tune_mappo.py                  | 94 +++++++++++++++----
 .../mappo_episode_metrics_callback.py         | 16 +++-
 .../training/train_mappo_bots.py              |  7 +-
 tests/test_mappo_episode_metrics_callback.py  | 30 ++++++
 5 files changed, 128 insertions(+), 20 deletions(-)
 create mode 100644 tests/test_mappo_episode_metrics_callback.py

diff --git a/configs/training/train_mappo_bots.yaml b/configs/training/train_mappo_bots.yaml
index 03aa8ef..5a94492 100644
--- a/configs/training/train_mappo_bots.yaml
+++ b/configs/training/train_mappo_bots.yaml
@@ -15,6 +15,7 @@ training:
   clip_range: 0.2
   ent_coef: 0.01
   vf_coef: 0.5
+  max_grad_norm: 0.5
   critic_hidden_dim: 256
   pool_max_size: 15
   checkpoint_interval: 100000
diff --git a/scripts/optuna_tune_mappo.py b/scripts/optuna_tune_mappo.py
index d69a972..6485b4b 100644
--- a/scripts/optuna_tune_mappo.py
+++ b/scripts/optuna_tune_mappo.py
@@ -14,7 +14,6 @@
 
 import optuna
 from loguru import logger
-from optuna import TrialPruned
 
 from village_ai_war.config_load import load_project_config
 from village_ai_war.training.train_mappo_bots import run_mappo_bots_training
@@ -28,6 +27,43 @@ def _study_path_slug(study_name: str) -> str:
 # Optuna requires the same categorical choices for a parameter name in every trial
 # (see CategoricalDistribution does not support dynamic value space).
 _BATCH_SIZE_CHOICES = [32, 64, 128, 256, 512, 1024]
+_N_STEPS_CHOICES = [128, 256, 512]
+
+# Correlated (gamma, gae_lambda) presets — single categorical for TPE-friendly search.
+_GAMMA_GAE_PRESETS = [
+    "0.97/0.92",
+    "0.98/0.94",
+    "0.99/0.95",
+    "0.995/0.97",
+    "0.999/0.98",
+]
+
+_CHECKPOINT_INTERVAL_CHOICES = [50_000, 100_000, 200_000]
+
+
+def _rollout_labels(n_envs: int) -> list[str]:
+    """Valid (n_steps, batch_size) pairs encoded as n{n}_bs{b} for fixed Optuna categorical space."""
+    labels: list[str] = []
+    for n_steps in _N_STEPS_CHOICES:
+        buf = n_envs * n_steps
+        for bs in _BATCH_SIZE_CHOICES:
+            if bs <= buf and buf % bs == 0:
+                labels.append(f"n{n_steps}_bs{bs}")
+    return sorted(labels)
+
+
+def _parse_rollout(label: str) -> tuple[int, int]:
+    if not label.startswith("n") or "_bs" not in label:
+        raise ValueError(f"bad rollout label: {label!r}")
+    left, right = label.split("_bs", 1)
+    return int(left[1:]), int(right)
+
+
+def _parse_gamma_gae(preset: str) -> tuple[float, float]:
+    parts = preset.split("/", 1)
+    if len(parts) != 2:
+        raise ValueError(f"bad gamma/gae preset: {preset!r}")
+    return float(parts[0]), float(parts[1])
 
 
 def main() -> None:
@@ -43,9 +79,9 @@ def main() -> None:
     )
     parser.add_argument("--seed", type=int, default=0)
     parser.add_argument("--config-name", type=str, default="default")
-    parser.add_argument("--n-envs", type=int, default=1)
-    parser.add_argument("--total-timesteps", type=int, default=50_000)
-    parser.add_argument("--selfplay-iterations", type=int, default=2)
+    parser.add_argument("--n-envs", type=int, default=4)
+    parser.add_argument("--total-timesteps", type=int, default=100_000)
+    parser.add_argument("--selfplay-iterations", type=int, default=4)
     parser.add_argument(
         "--disable-tensorboard",
         action="store_true",
@@ -60,27 +96,46 @@ def main() -> None:
     )
     args = parser.parse_args()
 
+    n_envs = int(args.n_envs)
+    rollout_labels = _rollout_labels(n_envs)
+    if not rollout_labels:
+        logger.error(
+            "No valid rollout (n_steps, batch_size) for n_envs={} with batch choices {}; "
+            "increase --n-envs or adjust _BATCH_SIZE_CHOICES.",
+            n_envs,
+            _BATCH_SIZE_CHOICES,
+        )
+        sys.exit(1)
+
     slug = _study_path_slug(args.study_name)
     user_overrides: list[str] = list(args.override)
 
+    logger.info(
+        "HPO search space uses rollout={}, gamma_gae presets={}, "
+        "checkpoint_intervals={}. If you changed param names vs an existing DB study, "
+        "use a new --study-name or storage file.",
+        len(rollout_labels),
+        len(_GAMMA_GAE_PRESETS),
+        _CHECKPOINT_INTERVAL_CHOICES,
+    )
+
     def objective(trial: optuna.Trial) -> float:
-        n_envs = int(args.n_envs)
-        n_steps = trial.suggest_categorical("n_steps", [128, 256, 512])
-        buffer = n_envs * n_steps
+        rollout = trial.suggest_categorical("rollout", rollout_labels)
+        n_steps, batch_size = _parse_rollout(rollout)
 
         lr = trial.suggest_float("learning_rate", 1e-5, 1e-2, log=True)
-        batch_size = trial.suggest_categorical("batch_size", _BATCH_SIZE_CHOICES)
-        if batch_size > buffer or buffer % batch_size != 0:
-            raise TrialPruned(
-                f"batch_size={batch_size} incompatible with n_envs*n_steps={buffer}"
-            )
         n_epochs = trial.suggest_int("n_epochs", 4, 15)
         ent_coef = trial.suggest_float("ent_coef", 1e-4, 0.1, log=True)
-        gamma = trial.suggest_float("gamma", 0.95, 0.999)
-        gae_lambda = trial.suggest_float("gae_lambda", 0.9, 0.99)
+        gamma_gae = trial.suggest_categorical("gamma_gae", _GAMMA_GAE_PRESETS)
+        gamma, gae_lambda = _parse_gamma_gae(gamma_gae)
         clip_range = trial.suggest_float("clip_range", 0.05, 0.3)
         vf_coef = trial.suggest_float("vf_coef", 0.1, 1.0)
         critic_hidden_dim = trial.suggest_categorical("critic_hidden_dim", [128, 256, 512])
+        max_grad_norm = trial.suggest_float("max_grad_norm", 0.3, 2.0)
+        pool_max_size = trial.suggest_int("pool_max_size", 8, 30)
+        checkpoint_interval = trial.suggest_categorical(
+            "checkpoint_interval", _CHECKPOINT_INTERVAL_CHOICES
+        )
 
         trial_dir = f"{slug}/trial_{trial.number}"
         overrides: list[str] = [
@@ -97,6 +152,9 @@ def objective(trial: optuna.Trial) -> float:
             f"training.clip_range={clip_range}",
             f"training.vf_coef={vf_coef}",
             f"training.critic_hidden_dim={critic_hidden_dim}",
+            f"training.max_grad_norm={max_grad_norm}",
+            f"training.pool_max_size={pool_max_size}",
+            f"training.checkpoint_interval={checkpoint_interval}",
             f"training.checkpoint_dir=checkpoints/optuna/{trial_dir}",
             f"training.pool_dir=checkpoints/optuna/{trial_dir}/pool",
             f"training.log_dir=logs/optuna/{trial_dir}",
@@ -108,14 +166,16 @@ def objective(trial: optuna.Trial) -> float:
         flat = load_project_config(_ROOT, config_name=args.config_name, overrides=overrides)
         metrics = run_mappo_bots_training(flat, return_metrics=True)
         assert metrics is not None
+        win_townhall_frac = float(metrics["win_townhall_frac"])
         win_frac = float(metrics["win_frac"])
         logger.info(
-            "trial {} finished: win_frac={} outcomes={}",
+            "trial {} finished: win_townhall_frac={} win_frac={} outcomes={}",
             trial.number,
+            win_townhall_frac,
             win_frac,
             metrics.get("outcome_fractions"),
         )
-        return win_frac
+        return win_townhall_frac
 
     storage = args.storage.strip() or None
     sampler = optuna.samplers.TPESampler(seed=args.seed)
@@ -129,7 +189,7 @@ def objective(trial: optuna.Trial) -> float:
     study.optimize(objective, n_trials=args.n_trials, n_jobs=args.n_jobs, show_progress_bar=True)
 
     best = study.best_trial
-    logger.info("Best trial: {} value={}", best.number, best.value)
+    logger.info("Best trial: {} win_townhall_frac={}", best.number, best.value)
     logger.info("Best params: {}", best.params)
 
 
diff --git a/src/village_ai_war/training/mappo_episode_metrics_callback.py b/src/village_ai_war/training/mappo_episode_metrics_callback.py
index 7185866..8dd8a44 100644
--- a/src/village_ai_war/training/mappo_episode_metrics_callback.py
+++ b/src/village_ai_war/training/mappo_episode_metrics_callback.py
@@ -15,6 +15,7 @@ def __init__(self, window: int = 512, verbose: int = 0) -> None:
         self._window = max(int(window), 1)
         self._outcomes: deque[str] = deque(maxlen=self._window)
         self._reasons: deque[str] = deque(maxlen=self._window)
+        self._win_townhall_flags: deque[bool] = deque(maxlen=self._window)
 
     def _on_step(self) -> bool:
         infos = self.locals.get("infos")
@@ -24,10 +25,14 @@ def _on_step(self) -> bool:
         for info in infos:
             if not isinstance(info, dict) or "episode_outcome" not in info:
                 continue
-            self._outcomes.append(str(info["episode_outcome"]))
+            outcome = str(info["episode_outcome"])
+            self._outcomes.append(outcome)
             tr = info.get("terminal_reason")
             if tr is not None:
                 self._reasons.append(str(tr))
+            self._win_townhall_flags.append(
+                outcome == "win" and info.get("terminal_reason") == "townhall_destroyed"
+            )
             updated = True
 
         if updated and self.logger is not None and self._outcomes:
@@ -35,6 +40,8 @@ def _on_step(self) -> bool:
             for k, v in Counter(self._outcomes).items():
                 self.logger.record(f"mappo/outcome_frac/{k}", float(v) / float(n))
             self.logger.record("mappo/episodes_in_window", float(n))
+            wth = sum(1 for x in self._win_townhall_flags if x) / float(n)
+            self.logger.record("mappo/win_townhall_frac", float(wth))
             rn = len(self._reasons)
             if rn > 0:
                 for k, v in Counter(self._reasons).items():
@@ -47,3 +54,10 @@ def outcome_fractions(self) -> dict[str, float]:
             return {}
         n = len(self._outcomes)
         return {str(k): float(v) / float(n) for k, v in Counter(self._outcomes).items()}
+
+    def win_townhall_frac(self) -> float:
+        """Fraction of episodes in the window that are a win via opponent town hall destroyed."""
+        if not self._win_townhall_flags:
+            return 0.0
+        n = len(self._win_townhall_flags)
+        return float(sum(1 for x in self._win_townhall_flags if x)) / float(n)
diff --git a/src/village_ai_war/training/train_mappo_bots.py b/src/village_ai_war/training/train_mappo_bots.py
index e8f9ba9..f29800d 100644
--- a/src/village_ai_war/training/train_mappo_bots.py
+++ b/src/village_ai_war/training/train_mappo_bots.py
@@ -41,8 +41,9 @@ def run_mappo_bots_training(
 ) -> dict[str, Any] | None:
     """MAPPO (PPO + centralized critic) with a rolling pool of past MAPPO snapshots.
 
-    If ``return_metrics`` is True, returns ``win_frac`` and ``outcome_fractions`` from the
-    episode metrics callback; otherwise returns ``None``.
+    If ``return_metrics`` is True, returns ``win_frac``, ``win_townhall_frac`` (fraction of
+    episodes in the metrics window that are wins via ``townhall_destroyed``), and
+    ``outcome_fractions`` from the episode metrics callback; otherwise returns ``None``.
     """
     flat = _flat_cfg(cfg)
     tcfg = flat["training"]
@@ -100,6 +101,7 @@ def _init() -> MAPPOBotEnv:
         clip_range=clip_range,
         ent_coef=ent_coef,
         vf_coef=vf_coef,
+        max_grad_norm=float(tcfg.get("max_grad_norm", 0.5)),
         tensorboard_log=tb_log,
         policy_kwargs={
             "map_size": n,
@@ -139,6 +141,7 @@ def _init() -> MAPPOBotEnv:
         fracs = metrics_cb.outcome_fractions()
         return {
             "win_frac": float(fracs.get("win", 0.0)),
+            "win_townhall_frac": float(metrics_cb.win_townhall_frac()),
             "outcome_fractions": fracs,
         }
     return None
diff --git a/tests/test_mappo_episode_metrics_callback.py b/tests/test_mappo_episode_metrics_callback.py
new file mode 100644
index 0000000..8eade2f
--- /dev/null
+++ b/tests/test_mappo_episode_metrics_callback.py
@@ -0,0 +1,30 @@
+"""Unit tests for MAPPOEpisodeMetricsCallback."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+
+pytest.importorskip("stable_baselines3")
+
+from village_ai_war.training.mappo_episode_metrics_callback import MAPPOEpisodeMetricsCallback
+
+
+def test_win_townhall_frac_counts_joint_outcome_and_reason() -> None:
+    cb = MAPPOEpisodeMetricsCallback(window=10, verbose=0)
+    cb.model = SimpleNamespace(logger=None)
+    cb.locals = {
+        "infos": [
+            {"episode_outcome": "win", "terminal_reason": "townhall_destroyed"},
+            {"episode_outcome": "win", "terminal_reason": "team1_eliminated"},
+            {"episode_outcome": "loss", "terminal_reason": "townhall_destroyed"},
+        ]
+    }
+    assert cb._on_step() is True
+    assert cb.win_townhall_frac() == pytest.approx(1.0 / 3.0)
+
+
+def test_win_townhall_frac_empty_window() -> None:
+    cb = MAPPOEpisodeMetricsCallback(window=10, verbose=0)
+    assert cb.win_townhall_frac() == 0.0

From d06af4cf44ed5bf65e7b7f2b50e29cd174e23f39 Mon Sep 17 00:00:00 2001
From: LizardAPN <piskunov-05@bk.ru>
Date: Wed, 1 Apr 2026 21:38:52 +0300
Subject: [PATCH 5/5] Add objective metric selection and mean episode reward
 tracking

- Introduced `--objective` argument in the Optuna tuning script to allow selection of metrics for optimization: `ep_rew_mean`, `win_townhall_frac`, and `win_frac`.
- Enhanced the episode metrics callback to compute and log the mean episode reward, improving performance tracking.
- Updated training function to return the mean episode reward alongside existing metrics, providing deeper insights into training outcomes.
- Added tests to validate the mean episode reward calculations, ensuring accuracy in metrics reporting.
---
 scripts/optuna_tune_mappo.py                  | 31 ++++++++++++++++---
 .../mappo_episode_metrics_callback.py         | 14 +++++++++
 .../training/train_mappo_bots.py              |  4 ++-
 tests/test_mappo_episode_metrics_callback.py  | 19 ++++++++++++
 4 files changed, 63 insertions(+), 5 deletions(-)

diff --git a/scripts/optuna_tune_mappo.py b/scripts/optuna_tune_mappo.py
index 6485b4b..01e6780 100644
--- a/scripts/optuna_tune_mappo.py
+++ b/scripts/optuna_tune_mappo.py
@@ -94,6 +94,17 @@ def main() -> None:
         metavar="KEY=VALUE",
         help="Extra Hydra-style overrides (repeatable). Applied after trial suggestions.",
     )
+    parser.add_argument(
+        "--objective",
+        type=str,
+        default="ep_rew_mean",
+        choices=("ep_rew_mean", "win_townhall_frac", "win_frac"),
+        help=(
+            "Metric to maximize. Use one objective per Optuna study (storage); "
+            "do not mix objectives when load_if_exists. "
+            "ep_rew_mean: dense signal from mean episode return in the metrics window."
+        ),
+    )
     args = parser.parse_args()
 
     n_envs = int(args.n_envs)
@@ -111,9 +122,10 @@ def main() -> None:
     user_overrides: list[str] = list(args.override)
 
     logger.info(
-        "HPO search space uses rollout={}, gamma_gae presets={}, "
+        "HPO objective={}; search space rollout={}, gamma_gae presets={}, "
         "checkpoint_intervals={}. If you changed param names vs an existing DB study, "
         "use a new --study-name or storage file.",
+        args.objective,
         len(rollout_labels),
         len(_GAMMA_GAE_PRESETS),
         _CHECKPOINT_INTERVAL_CHOICES,
@@ -168,14 +180,25 @@ def objective(trial: optuna.Trial) -> float:
         assert metrics is not None
         win_townhall_frac = float(metrics["win_townhall_frac"])
         win_frac = float(metrics["win_frac"])
+        mean_ep = float(metrics["mean_episode_reward"])
+        if args.objective == "ep_rew_mean":
+            value = mean_ep
+        elif args.objective == "win_townhall_frac":
+            value = win_townhall_frac
+        else:
+            value = win_frac
         logger.info(
-            "trial {} finished: win_townhall_frac={} win_frac={} outcomes={}",
+            "trial {} finished: objective={} value={} mean_episode_reward={} "
+            "win_townhall_frac={} win_frac={} outcomes={}",
             trial.number,
+            args.objective,
+            value,
+            mean_ep,
             win_townhall_frac,
             win_frac,
             metrics.get("outcome_fractions"),
         )
-        return win_townhall_frac
+        return value
 
     storage = args.storage.strip() or None
     sampler = optuna.samplers.TPESampler(seed=args.seed)
@@ -189,7 +212,7 @@ def objective(trial: optuna.Trial) -> float:
     study.optimize(objective, n_trials=args.n_trials, n_jobs=args.n_jobs, show_progress_bar=True)
 
     best = study.best_trial
-    logger.info("Best trial: {} win_townhall_frac={}", best.number, best.value)
+    logger.info("Best trial: {} {}={}", best.number, args.objective, best.value)
     logger.info("Best params: {}", best.params)
 
 
diff --git a/src/village_ai_war/training/mappo_episode_metrics_callback.py b/src/village_ai_war/training/mappo_episode_metrics_callback.py
index 8dd8a44..5a7de0c 100644
--- a/src/village_ai_war/training/mappo_episode_metrics_callback.py
+++ b/src/village_ai_war/training/mappo_episode_metrics_callback.py
@@ -16,6 +16,7 @@ def __init__(self, window: int = 512, verbose: int = 0) -> None:
         self._outcomes: deque[str] = deque(maxlen=self._window)
         self._reasons: deque[str] = deque(maxlen=self._window)
         self._win_townhall_flags: deque[bool] = deque(maxlen=self._window)
+        self._episode_returns: deque[float] = deque(maxlen=self._window)
 
     def _on_step(self) -> bool:
         infos = self.locals.get("infos")
@@ -33,6 +34,9 @@ def _on_step(self) -> bool:
             self._win_townhall_flags.append(
                 outcome == "win" and info.get("terminal_reason") == "townhall_destroyed"
             )
+            ep = info.get("episode")
+            if isinstance(ep, dict) and "r" in ep:
+                self._episode_returns.append(float(ep["r"]))
             updated = True
 
         if updated and self.logger is not None and self._outcomes:
@@ -42,6 +46,9 @@ def _on_step(self) -> bool:
             self.logger.record("mappo/episodes_in_window", float(n))
             wth = sum(1 for x in self._win_townhall_flags if x) / float(n)
             self.logger.record("mappo/win_townhall_frac", float(wth))
+            if self._episode_returns:
+                mer = sum(self._episode_returns) / float(len(self._episode_returns))
+                self.logger.record("mappo/mean_episode_reward", float(mer))
             rn = len(self._reasons)
             if rn > 0:
                 for k, v in Counter(self._reasons).items():
@@ -61,3 +68,10 @@ def win_townhall_frac(self) -> float:
             return 0.0
         n = len(self._win_townhall_flags)
         return float(sum(1 for x in self._win_townhall_flags if x)) / float(n)
+
+    def mean_episode_reward(self) -> float:
+        """Mean of episode returns in the rolling window (VecMonitor ``episode['r']``); 0 if none."""
+        if not self._episode_returns:
+            return 0.0
+        n = len(self._episode_returns)
+        return float(sum(self._episode_returns)) / float(n)
diff --git a/src/village_ai_war/training/train_mappo_bots.py b/src/village_ai_war/training/train_mappo_bots.py
index f29800d..86660af 100644
--- a/src/village_ai_war/training/train_mappo_bots.py
+++ b/src/village_ai_war/training/train_mappo_bots.py
@@ -42,7 +42,8 @@ def run_mappo_bots_training(
     """MAPPO (PPO + centralized critic) with a rolling pool of past MAPPO snapshots.
 
     If ``return_metrics`` is True, returns ``win_frac``, ``win_townhall_frac`` (fraction of
-    episodes in the metrics window that are wins via ``townhall_destroyed``), and
+    episodes in the metrics window that are wins via ``townhall_destroyed``),
+    ``mean_episode_reward`` (mean of VecMonitor episode returns in the same window), and
     ``outcome_fractions`` from the episode metrics callback; otherwise returns ``None``.
     """
     flat = _flat_cfg(cfg)
@@ -142,6 +143,7 @@ def _init() -> MAPPOBotEnv:
         return {
             "win_frac": float(fracs.get("win", 0.0)),
             "win_townhall_frac": float(metrics_cb.win_townhall_frac()),
+            "mean_episode_reward": float(metrics_cb.mean_episode_reward()),
             "outcome_fractions": fracs,
         }
     return None
diff --git a/tests/test_mappo_episode_metrics_callback.py b/tests/test_mappo_episode_metrics_callback.py
index 8eade2f..ff803f9 100644
--- a/tests/test_mappo_episode_metrics_callback.py
+++ b/tests/test_mappo_episode_metrics_callback.py
@@ -28,3 +28,22 @@ def test_win_townhall_frac_counts_joint_outcome_and_reason() -> None:
 def test_win_townhall_frac_empty_window() -> None:
     cb = MAPPOEpisodeMetricsCallback(window=10, verbose=0)
     assert cb.win_townhall_frac() == 0.0
+
+
+def test_mean_episode_reward_averages_episode_r() -> None:
+    cb = MAPPOEpisodeMetricsCallback(window=10, verbose=0)
+    cb.model = SimpleNamespace(logger=None)
+    cb.locals = {
+        "infos": [
+            {"episode_outcome": "win", "terminal_reason": "townhall_destroyed", "episode": {"r": 1.0}},
+            {"episode_outcome": "loss", "terminal_reason": "townhall_destroyed", "episode": {"r": 3.0}},
+            {"episode_outcome": "draw", "terminal_reason": "max_ticks", "episode": {"r": 5.0}},
+        ]
+    }
+    assert cb._on_step() is True
+    assert cb.mean_episode_reward() == pytest.approx(3.0)
+
+
+def test_mean_episode_reward_empty_window() -> None:
+    cb = MAPPOEpisodeMetricsCallback(window=10, verbose=0)
+    assert cb.mean_episode_reward() == 0.0