Update and restructure python

Author: Saurav Agarwal

VERSION

@@ -24,7 +24,9 @@ classifiers = [
 	"Environment :: GPU :: NVIDIA CUDA",
 ]
 dynamic = ["version"]
-dependencies = ["numpy"]
+dependencies = ["numpy", "pyyaml",
+  'toml; python_version < "3.11"',
+]
 
 [project.optional-dependencies]
 nn = ["torch >=2.1", "torchvision >=2.1"]

pyproject.toml

@@ -10,7 +10,9 @@
 
 from ._version import version as __version__
 from .core import *
+from .io_utils import IOUtils
+# from .nn import *
 
 __name__ = "coverage_control"
 
-__all__ = ["__version__", "core"]
+__all__ = ["__version__", "core", "nn", "IOUtils", "__name__"]

python/coverage_control/__init__.py

@@ -1,12 +1,10 @@
-
 from __future__ import annotations
 
 __name__ = "core"
 
-from .._core import Point2, PointVector, PolygonFeature, VoronoiCell, VoronoiCells
-from .._core import BivariateNormalDistribution, WorldIDF, RobotModel, CoverageSystem
+from .._core import Point2, PointVector, DblVector, DblVectorVector
+from .._core import PolygonFeature, VoronoiCell, VoronoiCells
 from .._core import Parameters
+from .._core import BivariateNormalDistribution, BNDVector, WorldIDF, RobotModel, CoverageSystem
 
-from .algorithms import *
-
-__all__ = ["Point2", "PointVector", "PolygonFeature", "VoronoiCell", "VoronoiCells", "BivariateNormalDistribution", "WorldIDF", "RobotModel", "CoverageSystem", "algorithms", "Parameters", "__name__"]
+__all__ = ["Point2", "PointVector", "DblVector", "DblVectorVector", "PolygonFeature", "VoronoiCell", "VoronoiCells", "BivariateNormalDistribution", "BNDVector", "WorldIDF", "RobotModel", "CoverageSystem", "Parameters", "__name__"]

python/coverage_control/core/__init__.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+__name__ = "nn"
+
+from .data_loaders import *
+from .models.cnn import CNN
+from .models.lpac import LPAC
+from .trainers import TrainModel
+
+__all__ = ["__name__", "DataLoaderUtils", "CoverageEnvUtils", "LocalMapCNNDataset", "LocalMapGNNDataset", "CNNGNNDataset", "VoronoiGNNDataset", "CNN", "LPAC", "TrainModel"]

python/coverage_control/core/algorithms/__init__.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+__name__ = "data_loaders"
+
+from .data_loader_utils import DataLoaderUtils
+from .loaders import LocalMapCNNDataset, LocalMapGNNDataset, CNNGNNDataset, VoronoiGNNDataset
+from .coverage_env_utils import CoverageEnvUtils
+
+__all__ = ["__name__", "DataLoaderUtils", "CoverageEnvUtils", "LocalMapCNNDataset", "LocalMapGNNDataset", "CNNGNNDataset", "VoronoiGNNDataset"]

python/coverage_control/nn/__init__.py

@@ -0,0 +1,328 @@
+#  This file is part of the CoverageControl library
+#
+#  Author: Saurav Agarwal
+#  Contact: [email protected], [email protected]
+#  Repository: https://github.com/KumarRobotics/CoverageControl
+#
+#  Copyright (c) 2024, Saurav Agarwal
+#
+#  The CoverageControl library is free software: you can redistribute it and/or
+#  modify it under the terms of the GNU General Public License as published by
+#  the Free Software Foundation, either version 3 of the License, or (at your
+#  option) any later version.
+#
+#  The CoverageControl library is distributed in the hope that it will be
+#  useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
+#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General
+#  Public License for more details.
+#
+#  You should have received a copy of the GNU General Public License along with
+#  CoverageControl library. If not, see <https://www.gnu.org/licenses/>.
+
+## @file coverage_env_utils.py
+#  @brief Utility functions for coverage environment
+
+import math
+import numpy
+import copy
+# import cv2
+import torch
+import torchvision
+import torch_geometric
+from torch_geometric.data import Dataset
+from scipy.ndimage import gaussian_filter
+from scipy.spatial import distance_matrix
+
+from ...core import PointVector, DblVector, DblVectorVector
+from ...core import Parameters
+from ...core import CoverageSystem
+
+## @ingroup python_api
+class CoverageEnvUtils:
+    """
+    Class for utility functions for coverage environment
+    """
+    @staticmethod
+    def to_tensor(data: object) -> torch.Tensor:
+        """
+        Converts various types of data to torch.Tensor
+
+        Can accept the following types:
+        - numpy.ndarray
+        - PointVector
+        - DblVectorVector
+        - DblVector
+
+        Args:
+            data: input data
+
+        Returns:
+            torch.Tensor: converted data
+
+        Raises:
+            ValueError: if data type is not supported
+
+        """
+        if isinstance(data, numpy.ndarray):
+            return torch.from_numpy(numpy.copy(data.astype(numpy.float32)))
+        elif isinstance(data, PointVector):
+            data_tensor = torch.Tensor(len(data), 2)
+            for i in range(len(data)):
+                data_tensor[i] = CoverageEnvUtils.to_tensor(data[i])
+            return data_tensor
+        elif isinstance(data, DblVectorVector):
+            data_tensor = torch.Tensor(len(data))
+            for i in range(len(data)):
+                data_tensor[i] = CoverageEnvUtils.to_tensor(data[i])
+            return data_tensor
+        elif isinstance(data, DblVector):
+            data_tensor = torch.Tensor(len(data))
+            for i in range(len(data)):
+                data_tensor[i] = float(data[i])
+            return data_tensor
+        else:
+            raise ValueError('Unknown data type: {}'.format(type(data)))
+
+    @staticmethod
+    def get_raw_local_maps(env: CoverageSystem, params: Parameters) -> torch.Tensor:
+        """
+        Get raw local maps
+
+        Args:
+            env: coverage environment
+            params: parameters
+
+        Returns:
+            torch.Tensor: raw local maps
+
+        """
+        local_maps = torch.zeros((env.GetNumRobots(), params.pLocalMapSize, params.pLocalMapSize))
+        for r_idx in range(env.GetNumRobots()):
+            local_maps[r_idx] = CoverageEnvUtils.to_tensor(env.GetRobotLocalMap(r_idx))
+        return local_maps
+
+    @staticmethod
+    def get_raw_obstacle_maps(env: CoverageSystem, params: Parameters) -> torch.Tensor:
+        """
+        Get raw obstacle maps
+
+        Args:
+            env: coverage environment
+            params: parameters
+
+        Returns:
+            torch.Tensor: raw obstacle maps
+
+        """
+        obstacle_maps = torch.zeros((env.GetNumRobots(), params.pLocalMapSize, params.pLocalMapSize))
+        for r_idx in range(env.GetNumRobots()):
+            obstacle_maps[r_idx] = CoverageEnvUtils.to_tensor(env.GetRobotObstacleMap(r_idx))
+        return obstacle_maps
+
+    @staticmethod
+    def get_communication_maps(env: CoverageSystem, params: Parameters, map_size: int) -> torch.Tensor:
+        """
+        Generate communication maps from positions
+
+        Communication maps are composed of two channels.
+        Each channnel has non-zero values for cells that correspond to the relative positions of the neighbors.
+        For the first channel, the value is the x-coordinate of the relative position divided by the communication range.
+        Similarly, the y-coordinte is used for the second channel.
+
+        Args:
+            env: coverage environment
+            params: parameters
+            map_size: size of the map
+
+        Returns:
+            torch.Tensor: communication maps
+        """
+        num_robots = env.GetNumRobots()
+
+        comm_maps = torch.zeros((num_robots, 2, map_size, map_size))
+        for r_idx in range(num_robots):
+            neighbors_pos = CoverageEnvUtils.to_tensor(env.GetRelativePositonsNeighbors(r_idx))
+            scaled_indices = torch.round(neighbors_pos * map_size / (params.pCommunicationRange * params.pResolution * 2.) + (map_size / 2. - params.pResolution / 2.))
+            # comm_range_mask = relative_dist[r_idx] < params.pCommunicationRange
+            # scaled_indices = scaled_relative_pos[r_idx][comm_range_mask]
+            indices = torch.transpose(scaled_indices, 1, 0)
+            indices = indices.long()
+            values = neighbors_pos / params.pCommunicationRange
+            # values = values / params.pCommunicationRange
+            # values = (values + params.pCommunicationRange) / (2. * params.pCommunicationRange)
+            comm_maps[r_idx][0] = torch.sparse_coo_tensor(indices, values[:, 0], torch.Size([map_size, map_size])).to_dense()
+            comm_maps[r_idx][1] = torch.sparse_coo_tensor(indices, values[:, 1], torch.Size([map_size, map_size])).to_dense()
+        return comm_maps
+        # positions = env.GetRobotPositions()
+        # robot_positions = CoverageEnvUtils.to_tensor(env.GetRobotPositions())
+        # relative_pos = robot_positions.unsqueeze(0) - robot_positions.unsqueeze(1)
+        # scaled_relative_pos = torch.round(relative_pos * map_size / (params.pCommunicationRange * params.pResolution * 2.) + (map_size / 2. - params.pResolution / 2.))
+        # relative_dist = relative_pos.norm(2, 2)
+        # diagonal_mask = torch.eye(num_robots).to(torch.bool)
+        # relative_dist.masked_fill_(diagonal_mask, params.pCommunicationRange + 1)
+
+    @staticmethod
+    def resize_maps(maps: torch.Tensor, resized_map_size: int) -> torch.Tensor:
+        """
+        Resize maps to a given size
+        Uses bilinear interpolation from torchvision.transforms.functional.resize
+        Options: antialias=True
+
+        Args:
+            maps: input maps
+            resized_map_size: size of the resized maps
+
+        Returns:
+            torch.Tensor: resized maps
+
+        """
+        shape = maps.shape
+        maps = maps.view(-1, maps.shape[-2], maps.shape[-1])
+        maps = torchvision.transforms.functional.resize(maps, (resized_map_size, resized_map_size), interpolation=torchvision.transforms.InterpolationMode.BILINEAR, antialias=True)
+        maps = maps.view(shape[:-2] + maps.shape[-2:])
+        return maps
+
+    @staticmethod
+    def get_maps(env: CoverageSystem, params: Parameters, resized_map_size: int, use_comm_map: bool) -> torch.Tensor:
+        """
+        Get maps for the coverage environment
+
+        Args:
+            env: coverage environment
+            params: parameters
+            resized_map_size: size of the resized maps
+            use_comm_map: whether to use communication maps
+
+        Returns:
+            torch.Tensor: maps
+
+        """
+
+        num_robots = env.GetNumRobots()
+        raw_local_maps = CoverageEnvUtils.get_raw_local_maps(env, params)
+        resized_local_maps = CoverageEnvUtils.resize_maps(raw_local_maps, resized_map_size)
+        raw_obstacle_maps = CoverageEnvUtils.get_raw_obstacle_maps(env, params)
+        resized_obstacle_maps = CoverageEnvUtils.resize_maps(raw_obstacle_maps, resized_map_size)
+        if use_comm_map:
+            comm_maps = CoverageEnvUtils.get_communication_maps(env, params, resized_map_size)
+            maps = torch.cat([resized_local_maps.unsqueeze(1), comm_maps, resized_obstacle_maps.unsqueeze(1)], 1)
+        else:
+            maps = torch.cat([resized_local_maps.unsqueeze(1), resized_obstacle_maps.unsqueeze(1)], 1)
+        return maps
+
+    @staticmethod
+    def get_voronoi_features(env: CoverageSystem) -> torch.Tensor:
+        """
+        Get voronoi features
+
+        Args:
+            env: coverage environment
+
+        Returns:
+            torch.Tensor: voronoi features
+        """
+        features = env.GetRobotVoronoiFeatures()
+        tensor_features = torch.zeros((len(features), len(features[0])))
+        for r_idx in range(len(features)):
+            tensor_features[r_idx] = CoverageEnvUtils.to_tensor(features[r_idx])
+        return tensor_features
+
+    @staticmethod
+    def get_robot_positions(env: CoverageSystem) -> torch.Tensor:
+        """
+        Get robot positions
+
+        Args:
+            env: coverage environment
+        
+        Returns:
+            torch.Tensor: robot positions
+        """
+        robot_positions = CoverageEnvUtils.to_tensor(env.GetRobotPositions())
+        return robot_positions
+
+    @staticmethod
+    def get_weights(env: CoverageSystem, params: Parameters) -> torch.Tensor:
+        onebyexp = 1. / math.exp(1.)
+        robot_positions = CoverageEnvUtils.to_tensor(env.GetRobotPositions())
+        pairwise_distances = torch.cdist(robot_positions, robot_positions, 2)
+        edge_weights = torch.exp(-(pairwise_distances.square())/(params.pCommunicationRange * params.pCommunicationRange))
+        edge_weights.masked_fill_(edge_weights < onebyexp, 0)
+        edge_weights.fill_diagonal_(0)
+        return edge_weights
+
+    # Legacy edge weights used in previous research
+    # The weights are proportional to the distance
+    # Trying to move away from this
+    @staticmethod
+    def robot_positions_to_edge_weights(robot_positions: PointVector, world_map_size: int, comm_range: float) -> torch.Tensor:
+        x = numpy.array(robot_positions)
+        S = distance_matrix(x, x)
+        S[S > comm_range] = 0
+        C = (world_map_size**2) / (S.shape[0]**2)
+        C = 3 / C
+        graph_obs = C * S
+        return graph_obs
+
+    @staticmethod
+    def get_torch_geometric_data(env: CoverageSystem, params: Parameters, use_cnn: bool, use_comm_map: bool, map_size: int) -> torch_geometric.data.Data:
+        """
+        Get torch geometric data
+        In this function, the edge weights are binary
+
+        Args:
+            env: coverage environment
+            params: parameters
+            use_cnn: whether to use CNN
+            use_comm_map: whether to use communication maps
+            map_size: size of the maps
+
+        Returns:
+            torch_geometric.data.Data: torch geometric data
+
+        """
+        if use_cnn:
+            features = CoverageEnvUtils.get_maps(env, params, map_size, use_comm_map)
+        else:
+            features = CoverageEnvUtils.get_voronoi_features(env)
+        edge_weights = CoverageEnvUtils.get_weights(env, params).to_sparse().coalesce()
+        edge_index = edge_weights.indices().long()
+        weights = edge_weights.values().float()
+        pos = CoverageEnvUtils.get_robot_positions(env)
+        pos = (pos + params.pWorldMapSize/2.0)/params.pWorldMapSize
+        data = torch_geometric.data.Data(
+                x=features,
+                edge_index=edge_index.clone().detach(),
+                edge_weight=weights.clone().detach(),
+                pos=pos.clone().detach()
+                )
+
+        return data
+
+    # Legacy maps which gives decent results
+    # Trying to move away from this
+    # @staticmethod
+    # def get_stable_maps(env, params, resized_map_size):
+    #     robot_positions = CoverageEnvUtils.to_tensor(env.GetRobotPositions())
+    #     num_robots = env.GetNumRobots()
+    #     maps = torch.empty((num_robots, 4, resized_map_size, resized_map_size))
+    #     h_vals = torch.linspace(1.0, -1.0, maps.shape[-2]+1)
+    #     h_vals = (h_vals[1:] + h_vals[:-1])/2
+    #     w_vals = torch.linspace(-1.0, 1.0, maps.shape[-1]+1)
+    #     w_vals = (w_vals[1:] + w_vals[:-1])/2
+    #     heatmap_x = torch.stack([h_vals] * maps.shape[-1], dim=1)/100
+    #     heatmap_y = torch.stack([w_vals] * maps.shape[-2], dim=0)/100
+    #     for r_idx in range(num_robots):
+    #         local_map = env.GetRobotLocalMap(r_idx)
+    #         resized_local_map = cv2.resize(local_map, dsize=(resized_map_size, resized_map_size), interpolation=cv2.INTER_AREA)
+    #         maps[r_idx][0] = torch.tensor(resized_local_map).float()
+
+    #         comm_map = env.GetCommunicationMap(r_idx)
+    #         filtered_comm_map = gaussian_filter(comm_map, sigma=(3,3), order=0)
+    #         resized_comm_map = torch.tensor(cv2.resize(numpy.array(filtered_comm_map), dsize=(resized_map_size, resized_map_size), interpolation=cv2.INTER_AREA)).float()
+    #         maps[r_idx][1] = resized_comm_map
+
+    #         maps[r_idx][2] = heatmap_x
+    #         maps[r_idx][3] = heatmap_y
+
+    #     return maps