"""Shared training hyperparameters for NephroSim environments.

These configs are the single source of truth for training hyperparameters.
Both the CLI entrypoint (agents/ppo.py) and train_all.py import from here.
"""

ENV_CONFIGS: dict[str, dict] = {
    "nephrosim/UreaClearing-v0": {
        "total_timesteps": 200_000,
        "learning_rate": 3e-4,
        "n_steps": 2048,
        "batch_size": 64,
        "gamma": 0.99,
        "gae_lambda": 0.95,
        "clip_range": 0.2,
        "ent_coef": 0.01,
        "seed": 42,
    },
    "nephrosim/UltrafiltrationControl-v0": {
        "total_timesteps": 200_000,
        "learning_rate": 3e-4,
        "n_steps": 2048,
        "batch_size": 64,
        "gamma": 0.99,
        "gae_lambda": 0.95,
        "clip_range": 0.2,
        "ent_coef": 0.01,
        "seed": 42,
    },
    "nephrosim/PhosphateMgmt-v0": {
        "total_timesteps": 200_000,
        "learning_rate": 1e-4,
        "n_steps": 1024,
        "batch_size": 64,
        "gamma": 0.99,
        "gae_lambda": 0.95,
        "clip_range": 0.2,
        "ent_coef": 0.01,
        "seed": 42,
    },
    "nephrosim/FullDialysisSession-v0": {
        "total_timesteps": 500_000,
        "learning_rate": 1e-4,
        "n_steps": 4096,
        "batch_size": 128,
        "gamma": 0.995,
        "gae_lambda": 0.95,
        "clip_range": 0.2,
        "ent_coef": 0.005,
        "seed": 42,
    },
}
