23Utility functions for coverage environment
37from scipy.spatial
import distance_matrix
39from ...core
import CoverageSystem, DblVector, DblVectorVector, Parameters, PointVector
43class CoverageEnvUtils:
45 Class for utility functions for coverage environment
49 def to_tensor(data: object) -> torch.Tensor:
51 Converts various types of data to torch.Tensor
53 Can accept the following types:
63 torch.Tensor: converted data
66 ValueError: if data type is not supported
70 if isinstance(data, numpy.ndarray):
71 return torch.from_numpy(numpy.copy(data.astype(numpy.float32)))
73 if isinstance(data, PointVector):
74 data_tensor = torch.Tensor(len(data), 2)
76 for i, _
in enumerate(data):
77 data_tensor[i] = CoverageEnvUtils.to_tensor(data[i])
81 if isinstance(data, DblVectorVector):
82 data_tensor = torch.Tensor(len(data))
84 for i, _
in enumerate(data):
85 data_tensor[i] = CoverageEnvUtils.to_tensor(data[i])
89 if isinstance(data, DblVector):
90 data_tensor = torch.Tensor(len(data))
92 for i, _
in enumerate(data):
93 data_tensor[i] = float(data[i])
96 raise ValueError(f
"Unknown data type: {type(data)}")
104 env: coverage environment
108 torch.Tensor: raw local maps
111 local_maps = torch.zeros(
112 (env.GetNumRobots(), params.pLocalMapSize, params.pLocalMapSize)
115 for r_idx
in range(env.GetNumRobots()):
116 local_maps[r_idx] = CoverageEnvUtils.to_tensor(env.GetRobotLocalMap(r_idx))
123 Get raw obstacle maps
126 env: coverage environment
130 torch.Tensor: raw obstacle maps
133 obstacle_maps = torch.zeros(
134 (env.GetNumRobots(), params.pLocalMapSize, params.pLocalMapSize)
137 for r_idx
in range(env.GetNumRobots()):
138 obstacle_maps[r_idx] = CoverageEnvUtils.to_tensor(
139 env.GetRobotObstacleMap(r_idx)
146 env: CoverageSystem, params: Parameters, map_size: int
149 Generate communication maps from positions
151 Communication maps are composed of two channels.
152 Each channnel has non-zero values for cells that correspond to the relative positions of the neighbors.
153 For the first channel, the value is the x-coordinate of the relative position divided by the communication range.
154 Similarly, the y-coordinte is used for the second channel.
157 env: coverage environment
159 map_size: size of the map
162 torch.Tensor: communication maps
164 num_robots = env.GetNumRobots()
166 comm_maps = torch.zeros((num_robots, 2, map_size, map_size))
168 for r_idx
in range(num_robots):
169 neighbors_pos = CoverageEnvUtils.to_tensor(
170 env.GetRelativePositonsNeighbors(r_idx)
172 scaled_indices = torch.round(
175 / (params.pCommunicationRange * params.pResolution * 2.0)
176 + (map_size / 2.0 - params.pResolution / 2.0)
180 indices = torch.transpose(scaled_indices, 1, 0)
181 indices = indices.long()
182 values = neighbors_pos / params.pCommunicationRange
185 comm_maps[r_idx][0] = torch.sparse_coo_tensor(
186 indices, values[:, 0], torch.Size([map_size, map_size])
188 comm_maps[r_idx][1] = torch.sparse_coo_tensor(
189 indices, values[:, 1], torch.Size([map_size, map_size])
202 def resize_maps(maps: torch.Tensor, resized_map_size: int) -> torch.Tensor:
204 Resize maps to a given size
205 Uses bilinear interpolation from torchvision.transforms.functional.resize
206 Options: antialias=True
210 resized_map_size: size of the resized maps
213 torch.Tensor: resized maps
217 maps = maps.view(-1, maps.shape[-2], maps.shape[-1])
218 maps = torchvision.transforms.functional.resize(
220 (resized_map_size, resized_map_size),
221 interpolation=torchvision.transforms.InterpolationMode.BILINEAR,
224 maps = maps.view(shape[:-2] + maps.shape[-2:])
232 resized_map_size: int,
236 Get maps for the coverage environment
239 env: coverage environment
241 resized_map_size: size of the resized maps
242 use_comm_map: whether to use communication maps
249 num_robots = env.GetNumRobots()
250 raw_local_maps = CoverageEnvUtils.get_raw_local_maps(env, params)
251 resized_local_maps = CoverageEnvUtils.resize_maps(
252 raw_local_maps, resized_map_size
254 raw_obstacle_maps = CoverageEnvUtils.get_raw_obstacle_maps(env, params)
255 resized_obstacle_maps = CoverageEnvUtils.resize_maps(
256 raw_obstacle_maps, resized_map_size
260 comm_maps = CoverageEnvUtils.get_communication_maps(
261 env, params, resized_map_size
265 resized_local_maps.unsqueeze(1),
267 resized_obstacle_maps.unsqueeze(1),
273 [resized_local_maps.unsqueeze(1), resized_obstacle_maps.unsqueeze(1)], 1
284 env: coverage environment
287 torch.Tensor: voronoi features
289 features = env.GetRobotVoronoiFeatures()
290 tensor_features = torch.zeros((len(features), len(features[0])))
292 for r_idx, _
in enumerate(features):
293 tensor_features[r_idx] = CoverageEnvUtils.to_tensor(features[r_idx])
295 return tensor_features
303 env: coverage environment
306 torch.Tensor: robot positions
308 robot_positions = CoverageEnvUtils.to_tensor(env.GetRobotPositions())
310 return robot_positions
313 def get_weights(env: CoverageSystem, params: Parameters) -> torch.Tensor:
315 Get edge weights for the communication graph
318 env: coverage environment
322 torch.Tensor: edge weights
324 onebyexp = 1.0 / math.exp(1.0)
325 robot_positions = CoverageEnvUtils.to_tensor(env.GetRobotPositions())
326 pairwise_distances = torch.cdist(robot_positions, robot_positions, 2)
327 edge_weights = torch.exp(
328 -(pairwise_distances.square())
329 / (params.pCommunicationRange * params.pCommunicationRange)
331 edge_weights.masked_fill_(edge_weights < onebyexp, 0)
332 edge_weights.fill_diagonal_(0)
341 robot_positions: PointVector, world_map_size: int, comm_range: float
344 Convert robot positions to edge weights
347 robot_positions: robot positions
348 world_map_size: size of the world map
349 comm_range: communication range
352 torch.Tensor: edge weights
354 x = numpy.array(robot_positions)
355 s_mat = distance_matrix(x, x)
356 s_mat[s_mat > comm_range] = 0
357 c_mat = (world_map_size**2) / (s_mat.shape[0] ** 2)
359 graph_obs = c_mat * s_mat
370 ) -> torch_geometric.data.Data:
372 Get torch geometric data
373 In this function, the edge weights are binary
376 env: coverage environment
378 use_cnn: whether to use CNN
379 use_comm_map: whether to use communication maps
380 map_size: size of the maps
383 torch_geometric.data.Data: torch geometric data
388 features = CoverageEnvUtils.get_maps(env, params, map_size, use_comm_map)
390 features = CoverageEnvUtils.get_voronoi_features(env)
391 edge_weights = CoverageEnvUtils.get_weights(env, params).to_sparse().coalesce()
392 edge_index = edge_weights.indices().long()
393 weights = edge_weights.values().float()
394 pos = CoverageEnvUtils.get_robot_positions(env)
395 pos = (pos + params.pWorldMapSize / 2.0) / params.pWorldMapSize
396 data = torch_geometric.data.Data(
398 edge_index=edge_index.clone().detach(),
399 edge_weight=weights.clone().detach(),
400 pos=pos.clone().detach(),
torch.Tensor get_weights(CoverageSystem env, Parameters params)
Get edge weights for the communication graph.
torch.Tensor resize_maps(torch.Tensor maps, int resized_map_size)
Resize maps to a given size Uses bilinear interpolation from torchvision.transforms....
torch.Tensor robot_positions_to_edge_weights(PointVector robot_positions, int world_map_size, float comm_range)
Convert robot positions to edge weights.
torch.Tensor get_voronoi_features(CoverageSystem env)
Get voronoi features.
torch.Tensor to_tensor(object data)
Converts various types of data to torch.Tensor.
torch.Tensor get_communication_maps(CoverageSystem env, Parameters params, int map_size)
Generate communication maps from positions.
torch_geometric.data.Data get_torch_geometric_data(CoverageSystem env, Parameters params, bool use_cnn, bool use_comm_map, int map_size)
Get torch geometric data In this function, the edge weights are binary.
torch.Tensor get_raw_local_maps(CoverageSystem env, Parameters params)
Get raw local maps.
torch.Tensor get_robot_positions(CoverageSystem env)
Get robot positions.
torch.Tensor get_raw_obstacle_maps(CoverageSystem env, Parameters params)
Get raw obstacle maps.
torch.Tensor get_maps(CoverageSystem env, Parameters params, int resized_map_size, bool use_comm_map)
Get maps for the coverage environment.