Feat new experiments and refactor

.gitignore

@@ -157,4 +157,7 @@ cython_debug/
 #  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
+#.idea/
+
+# Experiments results
+ray_results/

Dockerfile

@@ -0,0 +1,10 @@
+FROM rayproject/ray:2.38.0-py312-gpu
+
+RUN sudo apt update && sudo apt install -y libgdal-dev gdal-bin
+
+WORKDIR /app
+
+COPY pyproject.toml ./
+RUN pip install --no-cache-dir . && pip uninstall -y app
+
+CMD ["/bin/sh"]

README.md

@@ -9,6 +9,42 @@ Welcome to the official GitHub repository for the Drone Swarm Search (DSSE) algo
 
 Explore a diverse range of implementations that leverage the latest advancements in machine learning to solve complex coordination tasks in dynamic and unpredictable environments.
 
+## How to run
+
+Arguments inside [brackets] are optional.
+
+#### Run a training script
+
+```sh
+python run_script --file <training_file> [--exp_name <name of your experiment>] [--n_agents <number of agents>]
+```
+
+Example: training a MLP on the coverage environment.
+```sh
+python run_script --file train_ppo_mlp_cov.py --exp_name training_mlp_cov
+```
+
+
+#### Run a test script
+
+```sh
+python run_script.py --file <test script name> --checkpoint <path to checkpoint to evaluate> [--matrix_path <path to matrix if coverage env>] [--see] [--n_agents <number of agents>]
+```
+
+The --see switch makes shows you the agents playing and records a GIF on the env instead of collecting metrics.
+
+Example: Evaluating 4 agents trained on coverage env
+```sh
+python run_script.py --file test_trained_cov_mlp.py --checkpoint <my_checkpoint_path> --matrix_path data/min_matrix.npy --n_agents 4
+```
+
+
+#### Using the docker container
+
+```
+docker compose run --rm dsse-algorithms
+```
+
 ## 📚 Documentation Links
 
 - **[Documentation Site](https://pfeinsper.github.io/drone-swarm-search/)**: Access detailed tutorials, usage examples, and comprehensive technical documentation. This resource is essential for understanding the DSSE framework and integrating these algorithms into your projects effectively.

docker-compose.yml

@@ -0,0 +1,19 @@
+version: '3.9'
+
+services:
+  dsse-algorithms:
+    image: dsse-algorithms
+    build: .
+    container_name: dsse-algorithms
+    working_dir: /app
+    volumes:
+      - .:/app  # Mounts your current folder
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [gpu]
+    stdin_open: true
+    command: /bin/sh

pyproject.toml

@@ -0,0 +1,92 @@
+[project]
+name = "drone-swarm-search-algorithms"
+version = "0.1.0"
+description = "Algorithms to solve the DSSE environment, focusing on optimizing drone swarm search and navigation for critical applications."
+authors = [
+    { name="Ricardo Ribeiro Rodrigues", email="[email protected]" },
+    { name="Renato Laffranchi Falcão", email="[email protected]" },
+    { name="Pedro Henrique Britto Aragão Andrade", email="[email protected]" },
+    { name="Jorás Oliveira", email="[email protected]" },
+    { name="Fabricio Barth", email="[email protected]" },
+]
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "aiosignal==1.3.1",
+    "attrs==24.2.0",
+    "certifi==2024.8.30",
+    "charset-normalizer==3.3.2",
+    "click==8.1.7",
+    "cloudpickle==3.0.0",
+    "contourpy==1.3.0",
+    "cycler==0.12.1",
+    "dm-tree==0.1.8",
+    "dsse[coverage]==1.1.9",
+    "farama-notifications==0.0.4",
+    "filelock==3.16.1",
+    "fonttools==4.54.1",
+    "frozenlist==1.4.1",
+    "fsspec==2024.9.0",
+    "gymnasium==0.28.1",
+    "idna==3.10",
+    "imageio==2.35.1",
+    "jax-jumpy==1.0.0",
+    "jinja2==3.1.4",
+    "jsonschema-specifications==2023.12.1",
+    "jsonschema==4.23.0",
+    "kiwisolver==1.4.7",
+    "lazy-loader==0.4",
+    "llvmlite==0.43.0",
+    "lz4==4.3.3",
+    "markdown-it-py==3.0.0",
+    "markupsafe==2.1.5",
+    "matplotlib==3.8.4",
+    "mdurl==0.1.2",
+    "mpmath==1.3.0",
+    "msgpack==1.1.0",
+    "networkx==3.3",
+    "numba==0.60.0",
+    "numpy==2.0.2",
+    "nvidia-cublas-cu12==12.1.3.1",
+    "nvidia-cuda-cupti-cu12==12.1.105",
+    "nvidia-cuda-nvrtc-cu12==12.1.105",
+    "nvidia-cuda-runtime-cu12==12.1.105",
+    "nvidia-cudnn-cu12==9.1.0.70",
+    "nvidia-cufft-cu12==11.0.2.54",
+    "nvidia-curand-cu12==10.3.2.106",
+    "nvidia-cusolver-cu12==11.4.5.107",
+    "nvidia-cusparse-cu12==12.1.0.106",
+    "nvidia-nccl-cu12==2.20.5",
+    "nvidia-nvjitlink-cu12==12.6.68",
+    "nvidia-nvtx-cu12==12.1.105",
+    "packaging==24.1",
+    "pandas==2.2.3",
+    "pettingzoo==1.24.3",
+    "pillow==10.4.0",
+    "protobuf==5.28.2",
+    "pyarrow==17.0.0",
+    "pygame==2.6.1",
+    "pygments==2.18.0",
+    "pyparsing==3.1.4",
+    "python-dateutil==2.9.0.post0",
+    "pytz==2024.2",
+    "pyyaml==6.0.2",
+    "ray==2.38.0",
+    "referencing==0.35.1",
+    "requests==2.32.3",
+    "rich==13.8.1",
+    "rpds-py==0.20.0",
+    "scikit-image==0.24.0",
+    "scipy==1.14.1",
+    "shellingham==1.5.4",
+    "six==1.16.0",
+    "sympy==1.13.3",
+    "tensorboardx==2.6.2.2",
+    "tifffile==2024.9.20",
+    "torch==2.4.1",
+    "triton==3.0.0",
+    "typer==0.12.5",
+    "typing-extensions==4.12.2",
+    "tzdata==2024.2",
+    "urllib3==2.2.3",
+]

run_script.py

@@ -0,0 +1,68 @@
+import os
+import datetime
+from pathlib import Path
+import importlib
+import argparse
+
+argparser = argparse.ArgumentParser()
+argparser.add_argument("--file", type=str, required=True)
+argparser.add_argument("--checkpoint", type=str, required=False, default=None)
+argparser.add_argument(
+    "--matrix_path", type=str, required=False, default=argparse.SUPPRESS
+)
+argparser.add_argument("--see", action="store_true", default=False)
+argparser.add_argument(
+    "--storage_path",
+    type=str,
+    required=False,
+    default=f"{Path().resolve()}/ray_results/",
+)
+argparser.add_argument(
+    "--exp_name", type=str, required=False, default=str(datetime.datetime.now())
+)
+argparser.add_argument("--n_agents", type=int, required=False, default=2)
+argparser.add_argument("--use_random_positions", action="store_true", default=False)
+args = argparser.parse_args()
+
+
+def find_path_if_exists(file):
+    for root, _, files in os.walk("src"):
+        if file in files:
+            return os.path.join(root, file)
+    return None
+
+
+file = args.file
+path = find_path_if_exists(file)
+if path is None:
+    print(f"File {file} not found in src")
+    exit(1)
+
+all_modules_from_src = []
+all_dirs = path.split("/")
+for i, dir in enumerate(all_dirs):
+    if dir == "src":
+        break
+
+for dir in all_dirs[i:]:
+    all_modules_from_src.append(dir)
+
+if "test" in all_modules_from_src:
+    if args.checkpoint is None:
+        print(
+            "Checkpoint is required for test files, please provide it with --checkpoint"
+        )
+        exit(1)
+
+if "matrix_path" not in args:
+    args.matrix_path = "data/min_matrix.npy"
+    print(f"No --matrix_path provided, using default: {args.matrix_path}")
+
+full_path = ".".join(all_modules_from_src[:-1])
+module_path = f"{full_path}.{Path(file).stem}"
+module = importlib.import_module(module_path)
+
+if hasattr(module, "main"):
+    module.main(args)
+else:
+    print(f"The module does not have a 'main' function.")

src/coverage/test/test_trained_cov.py

@@ -0,0 +1,59 @@
+from DSSE import CoverageDroneSwarmSearch
+from DSSE.environment.wrappers import RetainDronePosWrapper, AllPositionsWrapper
+import ray
+from ray.rllib.env.wrappers.pettingzoo_env import ParallelPettingZooEnv
+from ray.rllib.models import ModelCatalog
+from ray.tune.registry import register_env
+from ray.rllib.algorithms.ppo import PPO
+from src.models.ppo_cnn import PpoCnnModel
+from src.models.cnn_config import CNNConfig
+from src.utils.play_env import play_with_record, evaluate_agent_coverage
+
+
+def main(args):
+    if args.matrix_path is None:
+        print("Please provide a matrix path")
+        exit(1)
+
+    # Register the model
+    model = PpoCnnModel
+    model.CONFIG = CNNConfig(kernel_sizes=[(3, 3), (2, 2)])
+    ModelCatalog.register_custom_model(model.NAME, model)
+
+    def env_creator(_):
+        print("-------------------------- ENV CREATOR --------------------------")
+        N_AGENTS = 2
+        env = CoverageDroneSwarmSearch(
+            timestep_limit=180,
+            drone_amount=N_AGENTS,
+            prob_matrix_path=args.matrix_path,
+            render_mode="human",
+        )
+        env = AllPositionsWrapper(env)
+        grid_size = env.grid_size
+        positions = [
+            (grid_size - 1, grid_size // 2),
+            (0, grid_size // 2),
+        ]
+        env = RetainDronePosWrapper(env, positions)
+        return env
+
+    env = env_creator(None)
+    register_env(
+        "DSSE_Coverage", lambda config: ParallelPettingZooEnv(env_creator(config))
+    )
+    ray.init()
+
+    checkpoint_path = args.checkpoint
+    PPOagent = PPO.from_checkpoint(checkpoint_path)
+
+    if args.see:
+        play_with_record(env, PPOagent)
+    else:
+        evaluate_agent_coverage(env, PPOagent)
+
+    env.close()
+
+
+if __name__ == "__main__":
+    main(None)

src/coverage/test/test_trained_cov_mlp.py

@@ -0,0 +1,54 @@
+from DSSE import CoverageDroneSwarmSearch
+from DSSE.environment.wrappers import RetainDronePosWrapper, AllFlattenWrapper
+import ray
+from ray.rllib.env.wrappers.pettingzoo_env import ParallelPettingZooEnv
+from ray.tune.registry import register_env
+from ray.rllib.algorithms.ppo import PPO
+from src.utils.play_env import play_with_record, evaluate_agent_coverage
+
+def position_on_cross(grid_size, n_agents):
+    positions = [
+        (0, grid_size // 2),
+        (grid_size - 1, grid_size // 2),
+        (grid_size // 2, 0),
+        (grid_size // 2, grid_size - 1),
+    ]
+    return positions[:n_agents]
+
+
+def main(args):
+    def env_creator(_):
+        print("-------------------------- ENV CREATOR --------------------------")
+        render_mode = "human" if args.see else "ansi"
+        # 6 hours of simulation, 600 radius
+        env = CoverageDroneSwarmSearch(
+            timestep_limit=200,
+            drone_amount=args.n_agents,
+            prob_matrix_path=args.matrix_path,
+            render_mode=render_mode,
+        )
+        env = AllFlattenWrapper(env)
+        grid_size = env.grid_size
+        print("Grid size: ", grid_size)
+        env = RetainDronePosWrapper(env, position_on_cross(grid_size, args.n_agents))
+        return env
+
+    env = env_creator(None)
+    register_env(
+        "DSSE_Coverage", lambda config: ParallelPettingZooEnv(env_creator(config))
+    )
+    ray.init()
+
+    checkpoint_path = args.checkpoint
+    PPOagent = PPO.from_checkpoint(checkpoint_path)
+
+    if args.see:
+        play_with_record(env, PPOagent)
+    else:
+        evaluate_agent_coverage(env, PPOagent)
+
+    env.close()
+
+
+if __name__ == "__main__":
+    main()