examples/rl_controller/train_ppo.py

"""Train a PPO controller for batch-size voltage regulation.

Trains one PPO model per datacenter site.  For multi-DC systems (ieee34,
ieee123), each site gets its own policy while other sites use fixed
mid-range batch sizes during that site's training.

Usage:
    python examples/rl_controller/train_ppo.py --system ieee13 \\
        --scenario-library examples/rl_controller/outputs/ieee13/scenario_library/train_n500 \\
        --total-timesteps 2000000
    python examples/rl_controller/train_ppo.py --system ieee13 \\
        --scenario-library .../train_n500 --obs-mode system-summary-only
    python examples/rl_controller/train_ppo.py --system ieee13 \\
        --scenario-library .../train_n500 --hidden-dims 256 256 256 --n-envs 8
"""

from __future__ import annotations

import csv
import logging
import sys
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import tyro
from env import (
    BatchSizeEnv,
    ObservationConfig,
    RewardConfig,
    ScenarioLibrary,
    SharedBatchSizeEnv,
    compute_bus_phase_groups,
    compute_zone_mask,
)
from scenarios import (
    EXPERIMENTS,
    DCSite,
    ScenarioOpenDSSGrid,
)

from openg2g.controller.tap_schedule import TapScheduleController
from openg2g.datacenter.config import (
    DatacenterConfig,
    InferenceModelSpec,
    ReplicaSchedule,
    TrainingRun,
)
from openg2g.datacenter.offline import OfflineDatacenter, OfflineWorkload
from openg2g.datacenter.workloads.inference import InferenceData
from openg2g.datacenter.workloads.training import TrainingTrace
from openg2g.grid.config import TapSchedule

from systems import (
    DT_CTRL,
    DT_DC,
    DT_GRID,
    POWER_AUG,
    SPECS_CACHE_DIR,
    TRAINING_TRACE_PATH,
    V_MAX,
    V_MIN,
)

logger = logging.getLogger(__name__)


def make_sim_factory(
    exp: dict,
    inference_data: InferenceData,
):
    """Return a callable that builds fresh simulation components.

    Returns `(make_sim, all_site_specs, all_replica_counts, all_initial_batch_sizes)`
    where `make_sim(scenario_override=None)` produces
    `(dict[str, DatacenterBackend], grid, tap_ctrl)`. When the env is sampling
    from a `ScenarioLibrary`, it passes the already-materialized scenario dict
    as `scenario_override`; otherwise `make_sim()` falls back to the
    experiment's defaults.
    """
    sys = exp["sys"]
    dc_sites: dict[str, DCSite] = exp["dc_sites"]
    pv_systems_base = exp.get("pv_systems", [])
    tvl_base = exp.get("time_varying_loads", [])
    training_base = exp.get("training_base")

    is_single_dc = len(dc_sites) == 1
    if is_single_dc:
        orig_sid = next(iter(dc_sites))
        dc_sites = {"_default": dc_sites[orig_sid]}

    all_site_specs: dict[str, tuple[InferenceModelSpec, ...]] = {}
    all_replica_counts: dict[str, dict[str, int]] = {}
    all_initial_batch_sizes: dict[str, dict[str, int]] = {}
    site_inference: dict[str, InferenceData] = {}
    for sid, site in dc_sites.items():
        specs = tuple(md.spec for md, _ in site.models)
        all_site_specs[sid] = specs
        all_replica_counts[sid] = {md.spec.model_label: sched.initial for md, sched in site.models}
        all_initial_batch_sizes[sid] = {md.spec.model_label: md.initial_batch_size for md, _ in site.models}
        site_inference[sid] = inference_data.filter_models(specs)

    _episode_counter = [0]

    def make_sim(scenario_override: dict | None = None):
        _episode_counter[0] += 1

        if scenario_override is not None:
            sites = scenario_override["dc_sites"]
            # Library was built keyed by the experiment's DC site id (e.g. "default");
            # the grid expects "_default" for single-DC. Remap once here.
            if is_single_dc and "_default" not in sites:
                orig = next(iter(sites))
                sites = {"_default": sites[orig]}
            pv_systems = scenario_override["pv_systems"]
            tvl = scenario_override["tvl"]
            training = scenario_override["training_run"]
        else:
            sites = dc_sites
            pv_systems = pv_systems_base
            tvl = tvl_base
            if training_base is not None:
                training = TrainingRun(
                    n_gpus=training_base["n_gpus"],
                    trace=training_base["trace"],
                    target_peak_W_per_gpu=training_base["target_peak_W_per_gpu"],
                ).at(t_start=training_base["t_start"], t_end=training_base["t_end"])
            else:
                training = None

        datacenters: dict[str, OfflineDatacenter] = {}
        for sid, site in sites.items():
            dc_config = DatacenterConfig(gpus_per_server=8, base_kw_per_phase=site.base_kw_per_phase)
            replica_schedules: dict[str, ReplicaSchedule] = {md.spec.model_label: sched for md, sched in site.models}
            initial_bs = {md.spec.model_label: md.initial_batch_size for md, _ in site.models}
            wl_kwargs: dict = {
                "inference_data": site_inference[sid],
                "replica_schedules": replica_schedules,
                "initial_batch_sizes": initial_bs,
            }
            if training is not None:
                wl_kwargs["training"] = training
            workload = OfflineWorkload(**wl_kwargs)
            datacenters[sid] = OfflineDatacenter(
                dc_config,
                workload,
                name=sid,
                dt_s=DT_DC,
                seed=site.seed,
                power_augmentation=POWER_AUG,
                total_gpu_capacity=site.total_gpu_capacity,
            )

        dc_config_pf = DatacenterConfig(base_kw_per_phase=0).power_factor
        exclude = tuple(sys.get("exclude_buses", ()))
        grid = ScenarioOpenDSSGrid(
            pv_systems=pv_systems,
            time_varying_loads=tvl,
            source_pu=sys["source_pu"],
            dss_case_dir=sys["dss_case_dir"],
            dss_master_file=sys["dss_master_file"],
            dt_s=DT_GRID,
            initial_tap_position=sys["initial_taps"],
            exclude_buses=exclude,
        )
        for sid, dc in datacenters.items():
            site = sites[sid]
            grid.attach_dc(
                dc,
                bus=site.bus,
                connection_type=site.connection_type,
                power_factor=dc_config_pf,
            )

        tap_ctrl = TapScheduleController(schedule=TapSchedule(()), dt_s=DT_CTRL)
        return datacenters, grid, tap_ctrl

    return make_sim, all_site_specs, all_replica_counts, all_initial_batch_sizes


def _new_episode_acc() -> dict:
    return {
        "voltage": 0.0,
        "throughput": 0.0,
        "latency": 0.0,
        "switch": 0.0,
        "safe": 0.0,
        "max_under": 0.0,
        "max_over": 0.0,
        "viol_frac_sum": 0.0,
        "n_steps": 0,
    }


class TrainingMetricsCallback:
    """SB3 BaseCallback that aggregates per-episode reward components and voltage stats.

    Writes one CSV row per completed episode and mirrors the same metrics to
    the SB3 TensorBoard logger so they show up alongside built-in PPO metrics.
    Imported lazily inside `main` so SB3 isn't a hard import for tooling that
    only wants the experiment definitions.
    """

    def __new__(cls, csv_path: Path):
        # Late-bind to BaseCallback so this module is importable without SB3.
        from stable_baselines3.common.callbacks import BaseCallback

        class _Impl(BaseCallback):
            def __init__(self, csv_path: Path):
                super().__init__(verbose=0)
                self.csv_path = csv_path
                self._per_env: dict[int, dict] = {}
                self._ep_count = 0
                self._fp = None
                self._writer = None

            def _on_training_start(self) -> None:
                self._fp = open(self.csv_path, "w", buffering=1, newline="")  # noqa: SIM115
                self._writer = csv.writer(self._fp)
                self._writer.writerow(
                    [
                        "episode",
                        "timestep",
                        "ep_reward",
                        "ep_length",
                        "voltage",
                        "throughput",
                        "latency",
                        "switch",
                        "safe",
                        "max_undervoltage",
                        "max_overvoltage",
                        "mean_violation_frac",
                    ]
                )

            def _on_step(self) -> bool:
                infos = self.locals.get("infos", []) or []
                for env_idx, info in enumerate(infos):
                    acc = self._per_env.setdefault(env_idx, _new_episode_acc())
                    rc = info.get("reward_components")
                    if rc is not None:
                        acc["voltage"] += rc.get("voltage", 0.0)
                        acc["throughput"] += rc.get("throughput", 0.0)
                        acc["latency"] += rc.get("latency", 0.0)
                        acc["switch"] += rc.get("switch", 0.0)
                        acc["safe"] += rc.get("safe", 0.0)
                    vs = info.get("voltage_stats")
                    if vs is not None:
                        if vs.get("max_under", 0.0) > acc["max_under"]:
                            acc["max_under"] = vs["max_under"]
                        if vs.get("max_over", 0.0) > acc["max_over"]:
                            acc["max_over"] = vs["max_over"]
                        acc["viol_frac_sum"] += vs.get("violation_frac", 0.0)
                    acc["n_steps"] += 1

                    # Monitor wrapper injects an "episode" key on done
                    ep = info.get("episode")
                    if ep is not None:
                        self._ep_count += 1
                        n = max(acc["n_steps"], 1)
                        row = [
                            self._ep_count,
                            self.num_timesteps,
                            float(ep["r"]),
                            int(ep["l"]),
                            acc["voltage"],
                            acc["throughput"],
                            acc["latency"],
                            acc["switch"],
                            acc["safe"],
                            acc["max_under"],
                            acc["max_over"],
                            acc["viol_frac_sum"] / n,
                        ]
                        self._writer.writerow(row)
                        # Mirror to TB
                        self.logger.record("custom/voltage_pen", acc["voltage"])
                        self.logger.record("custom/throughput_bonus", acc["throughput"])
                        self.logger.record("custom/latency_pen", acc["latency"])
                        self.logger.record("custom/switch_pen", acc["switch"])
                        self.logger.record("custom/safe_bonus", acc["safe"])
                        self.logger.record("custom/max_undervoltage", acc["max_under"])
                        self.logger.record("custom/max_overvoltage", acc["max_over"])
                        self.logger.record("custom/violation_frac", acc["viol_frac_sum"] / n)
                        self._per_env[env_idx] = _new_episode_acc()
                return True

            def _on_training_end(self) -> None:
                if self._fp is not None:
                    self._fp.close()
                    self._fp = None

        return _Impl(csv_path)


def plot_training_progress(csv_path: Path, output_path: Path, label: str) -> Path | None:
    """Read the per-episode metrics CSV and emit a 2x2 PNG dashboard.

    Returns the output path on success, or `None` if the CSV is empty.
    """
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    rows: list[dict] = []
    with open(csv_path, newline="") as fp:
        for r in csv.DictReader(fp):
            if any(v is None for v in r.values()):
                continue  # skip partial rows from interrupted buffered writes
            rows.append({k: float(v) if k not in ("episode", "ep_length") else int(float(v)) for k, v in r.items()})
    if not rows:
        return None

    eps = np.array([r["episode"] for r in rows])
    ep_reward = np.array([r["ep_reward"] for r in rows])
    voltage = np.array([r["voltage"] for r in rows])
    throughput = np.array([r["throughput"] for r in rows])
    latency = np.array([r["latency"] for r in rows])
    switch = np.array([r["switch"] for r in rows])
    max_under = np.array([r["max_undervoltage"] for r in rows])
    max_over = np.array([r["max_overvoltage"] for r in rows])
    viol_frac = np.array([r["mean_violation_frac"] for r in rows])

    def smooth(arr: np.ndarray, window: int) -> np.ndarray:
        if window <= 1 or len(arr) < 2:
            return arr
        kernel = np.ones(window) / window
        return np.convolve(arr, kernel, mode="same")

    window = max(1, len(rows) // 20)

    fig, axes = plt.subplots(2, 2, figsize=(12, 9))

    ax = axes[0, 0]
    ax.plot(eps, ep_reward, alpha=0.3, label="raw")
    ax.plot(eps, smooth(ep_reward, window), label=f"smooth (w={window})", linewidth=2)
    ax.set_xlabel("Episode")
    ax.set_ylabel("Episode reward")
    ax.set_title(f"Learning curve: {label}")
    ax.legend()
    ax.grid(alpha=0.3)

    ax = axes[0, 1]
    ax.plot(eps, smooth(voltage, window), label="voltage", color="C3")
    ax.plot(eps, smooth(throughput, window), label="throughput", color="C2")
    ax.plot(eps, smooth(latency, window), label="latency", color="C1")
    ax.plot(eps, smooth(switch, window), label="switch", color="C0")
    ax.set_xlabel("Episode")
    ax.set_ylabel("Component reward (per episode)")
    ax.set_title("Reward decomposition")
    ax.axhline(0, color="k", linewidth=0.5)
    ax.legend()
    ax.grid(alpha=0.3)

    ax = axes[1, 0]
    ax.plot(eps, smooth(max_under, window), label="max undervoltage", color="C0")
    ax.plot(eps, smooth(max_over, window), label="max overvoltage", color="C3")
    ax.set_xlabel("Episode")
    ax.set_ylabel("Worst per-step deviation (pu)")
    ax.set_title("Voltage violation magnitude")
    ax.legend()
    ax.grid(alpha=0.3)

    ax = axes[1, 1]
    ax.plot(eps, smooth(viol_frac, window), color="C4")
    ax.set_xlabel("Episode")
    ax.set_ylabel("Mean fraction of bus-phases violating")
    ax.set_title("Violation prevalence")
    ax.set_ylim(0, max(0.05, float(viol_frac.max()) * 1.1))
    ax.grid(alpha=0.3)

    fig.suptitle(f"PPO training progress: {label}", fontsize=14)
    fig.tight_layout()
    fig.savefig(output_path, dpi=110)
    plt.close(fig)
    return output_path


@dataclass
class Args:
    system: str = "ieee13"
    """System name (ieee13, ieee34, ieee123)."""
    total_timesteps: int = 200_000
    """Total environment timesteps for training (per site). Counted across ALL parallel envs."""
    learning_rate: float = 1e-4
    """PPO learning rate (initial value if lr_schedule != 'constant')."""
    lr_schedule: str = "constant"
    """Learning rate schedule: 'constant' or 'linear' (decays to 0 over training)."""
    n_steps: int = 3600
    """Rollout length per environment (one full simulated hour)."""
    checkpoint_interval_rollouts: int = 10
    """Checkpoint cadence in rollouts. With the defaults (n_steps=3600, n_envs=8), 10 rollouts ≈ 288000 env transitions, so a 2M-step run produces ~7 snapshots."""  # noqa: E501
    batch_size: int = 128
    """Minibatch size for PPO updates."""
    n_epochs: int = 10
    """Number of PPO epochs per update."""
    gamma: float = 0.99
    """Discount factor."""
    gae_lambda: float = 0.95
    """GAE lambda."""
    clip_range: float = 0.2
    """PPO clipping range."""
    ent_coef: float = 0.01
    """Entropy coefficient."""
    hidden_dims: tuple[int, ...] = (128, 128)
    """Hidden layer widths for the MLP policy/value network. Pass multiple values for a deeper net, e.g. --hidden-dims 256 256 256."""  # noqa: E501
    w_voltage: float = 1000.0
    """Reward weight for voltage violations."""
    w_throughput: float = 0.0
    """Reward weight for throughput. Default 0 to isolate the voltage-control objective."""
    w_latency: float = 0.0
    """Reward weight for latency violations. Default 0 to isolate the voltage-control objective."""
    w_switch: float = 0.01
    """Reward weight for switching cost (penalizes |log2(batch_t) - log2(batch_{t-1})| summed over models). Without this, randomized-scenario runs converge to a near-uniform action distribution and the deterministic eval policy ends up flipping batch sizes on every step. 0.01 is a gentle prior: much smaller than voltage_pen so it acts as a tie-breaker, not a co-equal objective."""  # noqa: E501
    w_safe: float = 0.0
    """Small positive reward for staying in the safe voltage range. Each step adds +w_safe * (fraction of bus-phases within [v_min, v_max]). Default 0 (disabled). Recommended: 0.01."""  # noqa: E501
    switch_mode: str = "magnitude"
    """Switch penalty mode: 'magnitude' (original log-ratio), 'binary' (fixed cost per change), or 'cooldown' (decaying cost, recent changes expensive)."""  # noqa: E501
    switch_cooldown_tau: float = 30.0
    """Time constant (steps) for cooldown switch penalty. Only used with --switch-mode cooldown."""
    action_mode: str = "delta"
    """Action space mode: 'delta' (per-model {-1,0,+1}, 3^N actions) or
    'coupled' (all models move by the same delta, 13 actions)."""
    reward_clip: float = 0.0
    """If > 0, clip per-step reward to [-reward_clip, +inf). Prevents catastrophic scenarios from dominating PPO updates. Recommended: 1.0 (affects ~4% of episodes, leaving normal training signal intact)."""  # noqa: E501
    vec_normalize: bool = True
    """Wrap the vec env with SB3 VecNormalize (running obs/reward normalization). Strongly recommended: voltage_pen variance across scenarios is huge and tanks value-function learning without it."""  # noqa: E501
    obs_mode: str = "full-voltage"
    """Voltage observation mode. Choices:
    - "full-voltage": all bus-phase raw voltages + per-system summary (3 global scalars).
    - "per-bus-summary": per-bus [min,max] phase voltage + per-zone-summary (if zones exist) or per-system summary.
    - "per-zone-summary": per-zone summary only (3 scalars/zone, no raw voltages). Requires zones.
    - "system-summary-only": 3 global scalars only (no raw voltages, no zone breakdown).
    """
    shared: bool = True
    """Train one shared PPO for all sites (instead of separate per-site)."""
    total_duration_s: int = 3600
    """Episode length in simulated seconds. Lower for fast smoke tests (e.g. 300 = 5 simulated minutes)."""
    n_envs: int = 1
    """Number of parallel rollout environments. >1 uses SubprocVecEnv (each subprocess builds its own OpenDSS instance to avoid global-state conflicts)."""  # noqa: E501
    tensorboard: bool = True
    """Write TensorBoard logs to <output_dir>/tb. View with `tensorboard --logdir <output_dir>/tb`."""
    plot: bool = True
    """Generate matplotlib training-progress plots after each model finishes."""
    output_dir: str = ""
    """Output directory (default: outputs/<system>/ppo)."""
    log_level: str = "INFO"
    """Logging verbosity."""
    scenario_library: str = ""
    """Path to a scenario library directory built by build_library.py (containing metadata.json + traces.npz). When set, episodes are sampled from this library."""  # noqa: E501
    ofo_baseline: bool = False
    """Subtract the OFO oracle's per-step voltage penalty from PPO's reward (requires --scenario-library). Disabling gives the raw voltage penalty as reward."""  # noqa: E501
    truncate_episode: bool = True
    """Fast-forward past the initial quiet period and terminate after the last violation (requires --scenario-library with t_control_start/end). Disabling uses full 3600s episodes."""  # noqa: E501
    seed: int = 42
    """Random seed."""
    init_from: str = ""
    """Path to a PPO checkpoint .zip to warm-start from (e.g. ppo_1152000_steps.zip). If a sibling ppo_vecnormalize_<steps>.pkl exists and --vec-normalize is set, its stats are loaded too. Hyperparameters stored in the checkpoint (lr, ent_coef, clip_range, …) are preserved; pass CLI flags only to change the env-side reward weights."""  # noqa: E501


def main() -> None:
    args = tyro.cli(Args)

    logging.basicConfig(
        level=getattr(logging, args.log_level),
        format="%(levelname)s %(asctime)s [%(name)s:%(lineno)d] %(message)s",
        datefmt="%H:%M:%S",
    )
    logging.getLogger("openg2g.coordinator").setLevel(logging.WARNING)
    logging.getLogger("openg2g.datacenter").setLevel(logging.WARNING)
    logging.getLogger("openg2g.grid").setLevel(logging.WARNING)

    if args.system not in EXPERIMENTS:
        logger.error("Unknown system: %s. Available: %s", args.system, list(EXPERIMENTS.keys()))
        sys.exit(1)

    training_trace = TrainingTrace.ensure(TRAINING_TRACE_PATH)

    exp = EXPERIMENTS[args.system](training_trace)
    script_dir = Path(__file__).resolve().parent
    output_dir = script_dir / "outputs" / args.system / (args.output_dir or "ppo")
    output_dir.mkdir(parents=True, exist_ok=True)

    # Collect all model specs across sites
    dc_sites: dict[str, DCSite] = exp["dc_sites"]
    all_specs: list[InferenceModelSpec] = []
    for site in dc_sites.values():
        all_specs.extend(md.spec for md, _ in site.models)
    all_specs_tuple = tuple({s.model_label: s for s in all_specs}.values())

    # Load data via the per-spec content-addressed cache under SPECS_CACHE_DIR.
    # InferenceData.ensure regenerates only specs whose manifest is missing.
    logger.info("Loading data for %s...", args.system)
    inference_data = InferenceData.ensure(
        SPECS_CACHE_DIR,
        all_specs_tuple,
        plot=False,
        dt_s=float(DT_DC),
    )

    from openg2g.controller.ofo import LogisticModelStore

    logistic_models = LogisticModelStore.ensure(
        SPECS_CACHE_DIR,
        all_specs_tuple,
        plot=False,
    )

    scenario_lib = None
    if args.scenario_library:
        scenario_lib = ScenarioLibrary(args.scenario_library)
        logger.info(
            "Loaded scenario library with %d scenarios from %s (ofo_baseline=%s, truncate=%s)",
            len(scenario_lib),
            args.scenario_library,
            args.ofo_baseline,
            args.truncate_episode,
        )

    make_sim, all_site_specs, all_replica_counts, all_initial_batch_sizes = make_sim_factory(
        exp,
        inference_data,
    )

    # Probe grid for v_index and n_bus_phases
    probe_dcs, probe_grid, _ = make_sim()
    for dc in probe_dcs.values():
        dc.do_reset()
        dc.start()
    probe_grid.do_reset()
    probe_grid.start()
    v_index = probe_grid.v_index
    n_bus_phases_full = len(v_index)
    probe_grid.stop()
    for dc in probe_dcs.values():
        dc.stop()

    _VALID_OBS_MODES = {"full-voltage", "per-bus-summary", "per-zone-summary", "system-summary-only"}
    if args.obs_mode not in _VALID_OBS_MODES:
        raise ValueError(f"--obs-mode must be one of {sorted(_VALID_OBS_MODES)}, got {args.obs_mode!r}")

    if args.action_mode not in ("delta", "coupled"):
        raise ValueError(f"--action-mode must be 'delta' or 'coupled', got {args.action_mode!r}")

    # Zone info needed early for per-zone-summary validation
    zones: dict[str, list[str]] | None = exp.get("sys", {}).get("zones")

    if args.obs_mode == "per-zone-summary" and zones is None:
        raise ValueError("--obs-mode per-zone-summary requires the system to have zones defined (e.g. ieee123)")

    if args.obs_mode == "full-voltage":
        n_bus_phases = n_bus_phases_full
        bus_phase_groups = None
    elif args.obs_mode == "per-bus-summary":
        bus_phase_groups = compute_bus_phase_groups(v_index)
        n_bus_phases = 2 * len(bus_phase_groups)
    else:  # per-zone-summary or system-summary-only
        n_bus_phases = 0
        bus_phase_groups = None

    logger.info(
        "Grid has %d bus-phase pairs across %d buses; obs_mode=%s, n_bus_phases=%d",
        n_bus_phases_full,
        len(set(b for b, _ in v_index)),
        args.obs_mode,
        n_bus_phases,
    )

    reward_config = RewardConfig(
        w_voltage=args.w_voltage,
        w_throughput=args.w_throughput,
        w_latency=args.w_latency,
        w_switch=args.w_switch,
        w_safe=args.w_safe,
        v_min=V_MIN,
        v_max=V_MAX,
        reward_clip=args.reward_clip,
        switch_mode=args.switch_mode,
        switch_cooldown_tau=args.switch_cooldown_tau,
    )

    site_ids = list(all_site_specs.keys())

    from stable_baselines3 import PPO
    from stable_baselines3.common.callbacks import CallbackList, CheckpointCallback
    from stable_baselines3.common.monitor import Monitor
    from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv, VecNormalize

    def _train_and_save(env_factory, label: str, save_name: str) -> None:
        """Build a (possibly vectorized) env from `env_factory` and train one PPO model.

        `env_factory` is a zero-arg callable returning a fresh `BatchSizeEnv`
        (or subclass). It is invoked once per parallel environment, wrapped with
        `Monitor`, and stitched into a vec-env. `SubprocVecEnv` is used when
        `args.n_envs > 1` because each rollout needs its own OpenDSS instance
        (OpenDSS holds global state, so multiple envs in one process collide).
        """

        def _make_one():
            env = env_factory()
            return Monitor(env)

        if args.n_envs > 1:
            vec_env = SubprocVecEnv([_make_one for _ in range(args.n_envs)])
        else:
            vec_env = DummyVecEnv([_make_one])

        if args.vec_normalize:
            vn_init_ckpt = None
            if args.init_from:
                _p = Path(args.init_from)
                _cand = _p.with_name(_p.name.replace("ppo_", "ppo_vecnormalize_", 1).replace(".zip", ".pkl"))
                if _cand.exists():
                    vn_init_ckpt = _cand
            if vn_init_ckpt is not None:
                vec_env = VecNormalize.load(str(vn_init_ckpt), vec_env)
                vec_env.training = True
                vec_env.norm_reward = True
                logger.info("Loaded VecNormalize stats from %s", vn_init_ckpt)
            else:
                if args.init_from:
                    logger.warning(
                        "--init-from set but no VecNormalize sibling pkl found; starting VecNormalize stats fresh."
                    )
                vec_env = VecNormalize(
                    vec_env,
                    norm_obs=True,
                    norm_reward=True,
                    clip_obs=10.0,
                    clip_reward=10.0,
                    gamma=args.gamma,
                )

        obs_dim = int(vec_env.observation_space.shape[0])
        if hasattr(vec_env.action_space, "nvec"):
            n_act = int(len(vec_env.action_space.nvec))
        elif hasattr(vec_env.action_space, "n"):
            n_act = int(vec_env.action_space.n)
        else:
            n_act = int(vec_env.action_space.shape[0])
        logger.info("")
        logger.info("=" * 60)
        logger.info("Training '%s': obs_dim=%d, n_actions=%d", label, obs_dim, n_act)
        logger.info(
            "  shared=%s, n_envs=%d, hidden_dims=%s, vec_normalize=%s",
            args.shared,
            args.n_envs,
            tuple(args.hidden_dims),
            args.vec_normalize,
        )
        logger.info(
            "  reward weights: voltage=%s throughput=%s latency=%s switch=%s safe=%s  reward_clip=%s",
            args.w_voltage,
            args.w_throughput,
            args.w_latency,
            args.w_switch,
            args.w_safe,
            args.reward_clip,
        )
        logger.info("  ofo_baseline=%s", args.ofo_baseline)
        logger.info(
            "  switch_mode=%s  switch_cooldown_tau=%s  action_mode=%s",
            args.switch_mode,
            args.switch_cooldown_tau,
            args.action_mode,
        )
        logger.info("=" * 60)

        checkpoint_cb = CheckpointCallback(
            save_freq=max(args.n_steps * args.checkpoint_interval_rollouts, 1),
            save_path=str(output_dir / "checkpoints" / label),
            name_prefix="ppo",
            save_vecnormalize=args.vec_normalize,
        )
        metrics_csv = output_dir / f"metrics_{label}.csv"
        metrics_cb = TrainingMetricsCallback(metrics_csv)
        callbacks = CallbackList([checkpoint_cb, metrics_cb])

        tb_log = str(output_dir / "tb") if args.tensorboard else None

        if args.lr_schedule == "linear":
            _lr_init = float(args.learning_rate)

            def lr_arg(progress_remaining):
                return progress_remaining * _lr_init

        elif args.lr_schedule == "constant":
            lr_arg = args.learning_rate
        else:
            raise ValueError(f"Unknown --lr-schedule: {args.lr_schedule!r} (expected 'constant' or 'linear')")

        if args.init_from:
            model = PPO.load(
                args.init_from,
                env=vec_env,
                device="auto",
                tensorboard_log=tb_log,
            )
            model.set_env(vec_env)
            logger.info(
                "Warm-started PPO from %s (num_timesteps=%d).",
                args.init_from,
                getattr(model, "num_timesteps", 0),
            )
        else:
            model = PPO(
                "MlpPolicy",
                vec_env,
                learning_rate=lr_arg,
                n_steps=args.n_steps,
                batch_size=args.batch_size,
                n_epochs=args.n_epochs,
                gamma=args.gamma,
                gae_lambda=args.gae_lambda,
                clip_range=args.clip_range,
                ent_coef=args.ent_coef,
                verbose=1,
                seed=args.seed,
                tensorboard_log=tb_log,
                policy_kwargs=dict(net_arch=list(args.hidden_dims)),
            )
        model.learn(
            total_timesteps=args.total_timesteps,
            callback=callbacks,
            tb_log_name=label,
            reset_num_timesteps=not bool(args.init_from),
        )
        model_path = output_dir / save_name
        model.save(str(model_path))
        logger.info("Saved '%s' model to %s.zip", label, model_path)
        if args.vec_normalize:
            # VecNormalize running stats MUST be reloaded at inference time,
            # otherwise the policy sees unnormalized obs and acts nonsensically.
            vn_path = output_dir / f"{save_name}_vecnormalize.pkl"
            vec_env.save(str(vn_path))
            logger.info("Saved VecNormalize stats to %s", vn_path)
        vec_env.close()

        if args.plot:
            try:
                plot_path = plot_training_progress(metrics_csv, output_dir / f"training_progress_{label}.png", label)
                if plot_path is not None:
                    logger.info("Wrote training plot to %s", plot_path)
                else:
                    logger.warning("No metrics rows in %s: skipping plot", metrics_csv)
            except Exception as e:
                logger.warning("Plotting failed for '%s': %s", label, e)

    if args.shared and len(site_ids) > 1:
        # ── Shared multi-site PPO ──
        logger.info("Training SHARED PPO for %d sites: %s", len(site_ids), site_ids)

        site_model_mapping = {sid: [s.model_label for s in all_site_specs[sid]] for sid in site_ids}
        all_initial_bs_flat = {label: bs for sid in site_ids for label, bs in all_initial_batch_sizes[sid].items()}
        zone_summary = (
            {zname: tuple(zbuses) for zname, zbuses in zones.items()}
            if zones is not None and args.obs_mode in ("per-zone-summary", "per-bus-summary")
            else None
        )
        obs_config = ObservationConfig.from_multi_site(
            all_site_specs,
            all_replica_counts,
            n_bus_phases=n_bus_phases,
            initial_batch_sizes=all_initial_bs_flat,
            zone_summary=zone_summary,
            bus_phase_groups=bus_phase_groups,
            v_min=V_MIN,
            v_max=V_MAX,
        )

        def shared_env_factory():
            return SharedBatchSizeEnv(
                make_sim_fn=make_sim,
                obs_config=obs_config,
                site_model_mapping=site_model_mapping,
                reward_config=reward_config,
                action_mode=args.action_mode,
                logistic_models=logistic_models,
                dt_ctrl=DT_CTRL,
                total_duration_s=args.total_duration_s,
                scenario_library=scenario_lib,
                ofo_baseline=args.ofo_baseline and scenario_lib is not None,
                truncate_episode=args.truncate_episode and scenario_lib is not None,
            )

        _train_and_save(shared_env_factory, "shared", "ppo_model_shared")

    else:
        # ── Per-site PPO ──
        logger.info("Training %d separate PPO(s): %s", len(site_ids), site_ids)

        for sid in site_ids:
            specs = all_site_specs[sid]
            replica_counts = all_replica_counts[sid]

            # Zone-local voltage filtering for systems with zone definitions
            zone_buses = None
            n_bp = n_bus_phases
            if zones is not None and sid in zones and args.obs_mode == "full-voltage":
                zone_buses = tuple(zones[sid])
                zone_mask = compute_zone_mask(v_index, zone_buses)
                n_bp = int(np.sum(zone_mask))
                logger.info("Site '%s': using zone-local obs with %d/%d bus-phases", sid, n_bp, n_bus_phases_full)

            site_initial_bs = all_initial_batch_sizes[sid]
            obs_config = ObservationConfig.from_model_specs(
                specs,
                replica_counts,
                n_bus_phases=n_bp,
                initial_batch_sizes=site_initial_bs,
                zone_buses=zone_buses,
                v_min=V_MIN,
                v_max=V_MAX,
            )

            def site_env_factory(_obs_config=obs_config, _sid=sid, _lib=scenario_lib):
                return BatchSizeEnv(
                    make_sim_fn=make_sim,
                    obs_config=_obs_config,
                    agent_site_id=_sid,
                    reward_config=reward_config,
                    action_mode=args.action_mode,
                    logistic_models=logistic_models,
                    dt_ctrl=DT_CTRL,
                    total_duration_s=args.total_duration_s,
                    scenario_library=_lib,
                    ofo_baseline=args.ofo_baseline and _lib is not None,
                    truncate_episode=args.truncate_episode and _lib is not None,
                )

            _train_and_save(site_env_factory, sid, f"ppo_model_{sid}")

    # Single-site runs alias the per-site output to the canonical ppo_model.zip
    # path that the docs + evaluate.py default to.
    if len(site_ids) == 1:
        import shutil

        for suffix in (".zip", "_vecnormalize.pkl"):
            src = output_dir / f"ppo_model_{site_ids[0]}{suffix}"
            dst = output_dir / f"ppo_model{suffix}"
            if src.exists():
                shutil.copy2(src, dst)

    logger.info("All done. Models saved to %s", output_dir)


if __name__ == "__main__":
    main()