swmmEnv 0.1.0__tar.gz

swmmenv-0.1.0/PKG-INFO

@@ -0,0 +1,106 @@
+Metadata-Version: 2.4
+Name: swmmEnv
+Version: 0.1.0
+Summary: Multi-agent reinforcement learning environment for SWMM stormwater simulation
+Home-page: https://github.com/DujDDx/swmmEnv
+Author: dujddx
+Author-email: dujddx <dujddx@163.com>
+License: MIT
+Project-URL: Homepage, https://github.com/DujDDx/swmmEnv
+Project-URL: Documentation, https://github.com/DujDDx/swmmEnv#readme
+Project-URL: Repository, https://github.com/DujDDx/swmmEnv.git
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: pyswmm>=2.1.0
+Requires-Dist: pettingzoo>=1.24.0
+Requires-Dist: gymnasium>=0.29.0
+Requires-Dist: numpy>=1.24.0
+Requires-Dist: pyyaml>=6.0
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.0; extra == "dev"
+Provides-Extra: marl
+Requires-Dist: marllib; extra == "marl"
+Requires-Dist: ray[rllib]>=2.0; extra == "marl"
+Dynamic: author
+Dynamic: home-page
+Dynamic: requires-python
+
+# SWMMEnv
+
+Multi-agent reinforcement learning environment for SWMM (Storm Water Management Model) simulation.
+
+## Overview
+
+SWMMEnv provides a PettingZoo-compatible interface for training multi-agent reinforcement learning (MARL) algorithms on stormwater management simulations. It integrates:
+
+- **PySWMM**: SWMM simulation engine
+- **PettingZoo**: Multi-agent RL environment interface
+- **MARLlib**: Training framework (MAPPO, QMIX, etc.)
+
+## Features
+
+- Read standard SWMM `.inp` files and rainfall `.dat` files
+- Control pump stations, gates, and weirs
+- Global reward for coupled stormwater systems
+- Config-driven design for different SWMM models
+- Time synchronization between RL steps and SWMM simulation steps
+- State normalization for stable training
+
+## Installation
+
+```bash
+pip install -e .
+```
+
+## Quick Start
+
+```python
+from swmmEnv import SWMMParallelEnv, load_config
+
+# Load configuration
+config = load_config("config/example.yaml")
+
+# Create environment
+env = SWMMParallelEnv(config)
+
+# Reset and get initial observations
+observations, info = env.reset()
+
+# Take a step
+actions = {"pump_1": 0.8, "gate_1": 0.5}
+observations, rewards, terminations, truncations, infos = env.step(actions)
+
+# Close environment
+env.close()
+```
+
+## Configuration
+
+See `config/default_config.yaml` for configuration structure.
+
+## Project Structure
+
+```
+swmmEnv/
+├── swmmEnv/
+│   ├── sim/           # Simulation modules (engine, time_sync, normalizer, mapping)
+│   ├── envs/          # RL environments (core MDP + PettingZoo wrapper)
+│   ├── reward/        # Reward functions
+│   └── config/        # Configuration system
+├── tests/             # Unit tests
+├── examples/          # Example scripts
+└── data/              # Sample SWMM models
+```
+
+## License
+
+MIT License

swmmenv-0.1.0/README.md

@@ -0,0 +1,70 @@
+# SWMMEnv
+
+Multi-agent reinforcement learning environment for SWMM (Storm Water Management Model) simulation.
+
+## Overview
+
+SWMMEnv provides a PettingZoo-compatible interface for training multi-agent reinforcement learning (MARL) algorithms on stormwater management simulations. It integrates:
+
+- **PySWMM**: SWMM simulation engine
+- **PettingZoo**: Multi-agent RL environment interface
+- **MARLlib**: Training framework (MAPPO, QMIX, etc.)
+
+## Features
+
+- Read standard SWMM `.inp` files and rainfall `.dat` files
+- Control pump stations, gates, and weirs
+- Global reward for coupled stormwater systems
+- Config-driven design for different SWMM models
+- Time synchronization between RL steps and SWMM simulation steps
+- State normalization for stable training
+
+## Installation
+
+```bash
+pip install -e .
+```
+
+## Quick Start
+
+```python
+from swmmEnv import SWMMParallelEnv, load_config
+
+# Load configuration
+config = load_config("config/example.yaml")
+
+# Create environment
+env = SWMMParallelEnv(config)
+
+# Reset and get initial observations
+observations, info = env.reset()
+
+# Take a step
+actions = {"pump_1": 0.8, "gate_1": 0.5}
+observations, rewards, terminations, truncations, infos = env.step(actions)
+
+# Close environment
+env.close()
+```
+
+## Configuration
+
+See `config/default_config.yaml` for configuration structure.
+
+## Project Structure
+
+```
+swmmEnv/
+├── swmmEnv/
+│   ├── sim/           # Simulation modules (engine, time_sync, normalizer, mapping)
+│   ├── envs/          # RL environments (core MDP + PettingZoo wrapper)
+│   ├── reward/        # Reward functions
+│   └── config/        # Configuration system
+├── tests/             # Unit tests
+├── examples/          # Example scripts
+└── data/              # Sample SWMM models
+```
+
+## License
+
+MIT License

swmmenv-0.1.0/examples/example_config.yaml

@@ -0,0 +1,76 @@
+# SWMMEnv Example Configuration
+
+# Input files (update paths to your actual SWMM model)
+inp_file: "data/sample.inp"
+rain_file: null  # If null, rainfall defined in inp file
+
+# Agents definition
+# Each agent controls one SWMM element
+agents:
+  pump_1:
+    type: pump
+    link_id: P1
+    upstream_node: J1
+    downstream_node: J2
+  gate_1:
+    type: weir
+    link_id: W1
+    upstream_node: J3
+    downstream_node: J4
+
+# Nodes to observe for state information
+obs_nodes:
+  - J1
+  - J2
+  - J3
+  - J4
+
+# Raingage ID for rainfall observation
+obs_raingage: RG1
+
+# Time synchronization
+# RL agent makes decisions every decision_interval seconds
+# SWMM advances by swmm_step seconds per simulation step
+time_sync:
+  decision_interval: 300  # 5 minutes
+  swmm_step: 10           # 10 seconds
+
+# Observation/reward normalization parameters
+normalization:
+  obs:
+    depth:
+      mean: 2.0
+      std: 1.5
+    flow:
+      mean: 0.5
+      std: 0.3
+    rainfall:
+      mean: 5.0
+      std: 10.0
+    setting:
+      mean: 0.5
+      std: 0.3
+  reward:
+    mean: 0.0
+    std: 10.0
+
+# Reward function name (from swmmEnv.reward)
+reward_fn: "default_reward"
+
+# Reward weights for default_reward
+reward_weights:
+  flooding: 1.0
+  level_deviation: 0.5
+  energy: 0.2
+
+# Target water levels (optional, for level tracking reward)
+target_levels:
+  J1: 1.5
+  J2: 1.0
+
+# Episode configuration
+max_steps: 1000
+warmup_steps: 0  # Steps to run before episode starts
+
+# Hotstart file for consistent resets (optional)
+hotstart_file: null

swmmenv-0.1.0/examples/manual_control.py

@@ -0,0 +1,214 @@
+"""
+Manual control example for SWMMEnv.
+
+This script demonstrates how to manually control the SWMM environment
+without MARLlib, useful for testing and debugging.
+"""
+
+import sys
+import os
+
+# Add parent directory to path for imports
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from swmmEnv import SWMMParallelEnv, load_config
+
+
+def run_manual_control(config_path: str = None, num_episodes: int = 1):
+    """
+    Run manual control demonstration.
+
+    Args:
+        config_path: Path to configuration file
+        num_episodes: Number of episodes to run
+    """
+    # Load configuration
+    if config_path is None:
+        config = load_config(None)
+        print("Using default configuration")
+    else:
+        config = load_config(config_path)
+        print(f"Loaded configuration from: {config_path}")
+
+    # Create environment
+    env = SWMMParallelEnv(config)
+
+    print(f"\nEnvironment created:")
+    print(f"  Agents: {env.possible_agents}")
+    print(f"  Max steps: {config.get('max_steps', 1000)}")
+
+    for episode in range(num_episodes):
+        print(f"\n{'='*60}")
+        print(f"Episode {episode + 1}/{num_episodes}")
+        print('='*60)
+
+        # Reset environment
+        observations, info = env.reset()
+
+        print(f"\nInitial observations:")
+        for agent, obs in observations.items():
+            print(f"  {agent}: {obs}")
+
+        # Run episode with random actions
+        step = 0
+        total_reward = 0.0
+        done = False
+
+        while not done:
+            # Generate random actions
+            actions = {}
+            for agent in env.agents:
+                # Random action between 0 and 1
+                import numpy as np
+                actions[agent] = np.array([np.random.random()], dtype=np.float32)
+
+            # Step environment
+            observations, rewards, terminations, truncations, infos = env.step(actions)
+
+            # Track reward (all agents get same reward)
+            step_reward = list(rewards.values())[0]
+            total_reward += step_reward
+            step += 1
+
+            # Check done
+            done = any(terminations.values())
+
+            # Print progress every 10 steps
+            if step % 10 == 0:
+                print(f"  Step {step}: reward={step_reward:.3f}, "
+                      f"cumulative={total_reward:.3f}")
+
+            # Render every 50 steps
+            if step % 50 == 0:
+                env.render()
+
+        print(f"\nEpisode finished after {step} steps")
+        print(f"Total reward: {total_reward:.3f}")
+
+        # Get final state
+        state = env.core_env.get_state()
+        print(f"\nFinal state:")
+        print(f"  Total flooding: {state.get('rainfall', 0):.1f} mm/h")
+        print(f"  Rainfall: {env.core_env.engine.get_total_flooding():.3f} m³/s")
+
+    # Close environment
+    env.close()
+    print("\nEnvironment closed.")
+
+
+def interactive_control(config_path: str = None):
+    """
+    Interactive control mode.
+
+    Allows user to input actions manually.
+    """
+    config = load_config(config_path)
+    env = SWMMParallelEnv(config)
+
+    print("\n" + "="*60)
+    print("Interactive SWMM Control")
+    print("="*60)
+    print(f"Agents: {env.possible_agents}")
+    print("\nCommands:")
+    print("  Enter action value (0.0-1.0) for each agent")
+    print("  'q' - Quit")
+    print("  'r' - Reset episode")
+    print("  's' - Show state")
+    print("="*60)
+
+    observations, _ = env.reset()
+    step = 0
+    total_reward = 0.0
+
+    while True:
+        print(f"\n--- Step {step} ---")
+
+        # Get actions from user
+        actions = {}
+        for agent in env.agents:
+            while True:
+                try:
+                    value = input(f"  {agent} [0.0-1.0]: ").strip()
+
+                    if value.lower() == 'q':
+                        env.close()
+                        print("Goodbye!")
+                        return
+
+                    if value.lower() == 'r':
+                        observations, _ = env.reset()
+                        step = 0
+                        total_reward = 0.0
+                        print("Episode reset!")
+                        break
+
+                    if value.lower() == 's':
+                        env.render()
+                        continue
+
+                    action = float(value)
+                    action = max(0.0, min(1.0, action))
+                    import numpy as np
+                    actions[agent] = np.array([action], dtype=np.float32)
+                    break
+
+                except ValueError:
+                    print("  Invalid input. Enter a number between 0.0 and 1.0")
+
+        if not actions:
+            continue
+
+        # Step environment
+        observations, rewards, terminations, truncations, infos = env.step(actions)
+
+        step_reward = list(rewards.values())[0]
+        total_reward += step_reward
+        step += 1
+
+        print(f"  Reward: {step_reward:.3f} (cumulative: {total_reward:.3f})")
+
+        # Check termination
+        if any(terminations.values()):
+            print("\nEpisode ended!")
+            print(f"Total steps: {step}")
+            print(f"Total reward: {total_reward:.3f}")
+
+            response = input("\nStart new episode? (y/n): ").strip().lower()
+            if response == 'y':
+                observations, _ = env.reset()
+                step = 0
+                total_reward = 0.0
+            else:
+                break
+
+    env.close()
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="SWMMEnv manual control")
+    parser.add_argument(
+        '--config',
+        type=str,
+        default=None,
+        help='Path to configuration file'
+    )
+    parser.add_argument(
+        '--episodes',
+        type=int,
+        default=1,
+        help='Number of episodes to run'
+    )
+    parser.add_argument(
+        '--interactive',
+        action='store_true',
+        help='Run in interactive mode'
+    )
+
+    args = parser.parse_args()
+
+    if args.interactive:
+        interactive_control(args.config)
+    else:
+        run_manual_control(args.config, args.episodes)

swmmenv-0.1.0/examples/train_mappo.py

@@ -0,0 +1,197 @@
+"""
+MARLlib training example for SWMMEnv.
+
+This script demonstrates how to train a multi-agent policy
+using MARLlib with the SWMMEnv.
+"""
+
+import sys
+import os
+
+# Add parent directory to path for imports
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+
+def train_mappo(config_path: str = None, num_steps: int = 100000):
+    """
+    Train MAPPO on SWMMEnv using MARLlib.
+
+    Args:
+        config_path: Path to SWMMEnv configuration file
+        num_steps: Number of training timesteps
+    """
+    try:
+        from marllib import marl
+        import ray
+    except ImportError:
+        print("MARLlib is required for this example.")
+        print("Install with: pip install marllib ray[rllib]")
+        return
+
+    # Import environment registration
+    from swmmEnv.envs.register_env import make_env, register_with_marllib
+
+    # Register environment with MARLlib
+    register_with_marllib()
+
+    # Create environment
+    if config_path:
+        env = make_env(
+            environment_name="swmm",
+            map_name="custom",
+            config_path=config_path
+        )
+    else:
+        env = make_env(environment_name="swmm", map_name="control")
+
+    # Get environment info
+    env_info = env.get_env_info()
+    print("\nEnvironment info:")
+    print(f"  Number of agents: {env_info['num_agents']}")
+    print(f"  Episode limit: {env_info['episode_limit']}")
+
+    # Build model configuration
+    model = marl.build_model(
+        environment=env,
+        algorithm=marl.algos.mappo,
+        model_preference={
+            "core_arch": "mlp",  # MLP for non-sequential control
+            "encode_layer": "128-128",
+            "hidden_dim": 64,
+        }
+    )
+
+    # Initialize MAPPO algorithm
+    mappo = marl.algos.mappo(
+        hyperparam_source="common"
+    )
+
+    # Training configuration
+    train_config = {
+        "lr": 0.0005,
+        "gamma": 0.99,
+        "batch_episode": 10,
+        "num_sgd_iter": 5,
+        "vf_loss_coeff": 1.0,
+        "entropy_coeff": 0.01,
+        "clip_param": 0.3,
+    }
+
+    print("\nStarting training...")
+    print(f"  Algorithm: MAPPO")
+    print(f"  Target timesteps: {num_steps}")
+
+    # Start training
+    mappo.fit(
+        env=env,
+        model=model,
+        stop={"timesteps_total": num_steps},
+        **train_config
+    )
+
+    print("\nTraining completed!")
+
+
+def train_with_rllib(config_path: str = None, num_steps: int = 100000):
+    """
+    Train using Ray RLlib directly (alternative to MARLlib).
+
+    This shows direct RLlib integration without MARLlib's abstraction.
+    """
+    try:
+        from ray.rllib.algorithms.ppo import PPOConfig
+        from ray.rllib.env import ParallelPettingZooEnv
+        import ray
+    except ImportError:
+        print("Ray RLlib is required for this example.")
+        print("Install with: pip install ray[rllib]")
+        return
+
+    from swmmEnv import SWMMParallelEnv, load_config
+
+    # Load config
+    config = load_config(config_path)
+
+    # Initialize Ray
+    ray.init(ignore_reinit_error=True)
+
+    # Create environment creator
+    def env_creator(env_config):
+        env_config_inner = config.copy()
+        env_config_inner.update(env_config)
+        return SWMMParallelEnv(env_config_inner)
+
+    # Register environment with Ray
+    from ray.tune.registry import register_env
+    register_env("swmm_env", env_creator)
+
+    # Configure PPO for multi-agent
+    config = (
+        PPOConfig()
+        .environment("swmm_env")
+        .multi_agent(
+            policies={
+                "shared_policy": None,  # Will be auto-generated
+            },
+            policy_mapping_fn=lambda agent_id: "shared_policy",
+        )
+        .training(
+            lr=0.0005,
+            gamma=0.99,
+            train_batch_size=4000,
+            sgd_minibatch_size=128,
+            num_sgd_iter=10,
+        )
+        .resources(
+            num_gpus=0,  # Set to 1 if GPU available
+        )
+    )
+
+    # Build algorithm
+    algo = config.build()
+
+    print("\nStarting training with Ray RLlib...")
+
+    # Training loop
+    for i in range(num_steps // 4000):
+        result = algo.train()
+        print(f"Iteration {i}: "
+              f"timesteps={result['timesteps_total']}, "
+              f"episode_reward_mean={result['episode_reward_mean']:.2f}")
+
+    print("\nTraining completed!")
+
+    # Cleanup
+    ray.shutdown()
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Train MARL on SWMMEnv")
+    parser.add_argument(
+        '--config',
+        type=str,
+        default=None,
+        help='Path to SWMMEnv configuration file'
+    )
+    parser.add_argument(
+        '--steps',
+        type=int,
+        default=100000,
+        help='Number of training timesteps'
+    )
+    parser.add_argument(
+        '--backend',
+        type=str,
+        choices=['marllib', 'rllib'],
+        default='marllib',
+        help='Training backend to use'
+    )
+
+    args = parser.parse_args()
+
+    if args.backend == 'marllib':
+        train_mappo(args.config, args.steps)
+    else:
+        train_with_rllib(args.config, args.steps)

swmmenv-0.1.0/pyproject.toml

@@ -0,0 +1,51 @@
+[build-system]
+requires = ["setuptools>=45.0", "wheel", "setuptools-scm>=6.2"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "swmmEnv"
+version = "0.1.0"
+description = "Multi-agent reinforcement learning environment for SWMM stormwater simulation"
+readme = "README.md"
+license = {text = "MIT"}
+authors = [
+    {name = "dujddx", email = "dujddx@163.com"}
+]
+requires-python = ">=3.8"
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.8",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+]
+
+dependencies = [
+    "pyswmm>=2.1.0",
+    "pettingzoo>=1.24.0",
+    "gymnasium>=0.29.0",
+    "numpy>=1.24.0",
+    "pyyaml>=6.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=7.0",
+    "pytest-cov>=4.0",
+]
+marl = [
+    "marllib",
+    "ray[rllib]>=2.0",
+]
+
+[project.urls]
+Homepage = "https://github.com/DujDDx/swmmEnv"
+Documentation = "https://github.com/DujDDx/swmmEnv#readme"
+Repository = "https://github.com/DujDDx/swmmEnv.git"
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["swmmEnv*"]