Skip to content

Commit 98872d5

Browse files
author
Saurav Agarwal
committed
Update data storage
1 parent ce1ced3 commit 98872d5

8 files changed

Lines changed: 201 additions & 212 deletions

File tree

‎cppsrc/main/data_generation.cpp‎

Lines changed: 6 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -20,72 +20,13 @@ namespace CC = CoverageControl;
2020
namespace CCT = CoverageControlTorch;
2121

2222
int main(int argc, char** argv) {
23-
/* if (argc != 2) { */
24-
/* std::cout << "Usage: ./data_generation <path_to_resizer_model>" << std::endl; */
25-
/* std::cout << "TorchVision JIT Transformer is required\n"; */
26-
/* return 1; */
27-
/* } */
28-
29-
/* // Read first argument as path to resizer model */
30-
/* std::string resizer_model_path = argv[1]; */
31-
/* if (!std::filesystem::exists(resizer_model_path)) { */
32-
/* std::cout << "Error: " << resizer_model_path << " does not exist\n"; */
33-
/* return 1; */
34-
/* } */
35-
36-
/* Parameters params; */
37-
38-
/* std::vector <Point2> robot_positions; */
39-
/* robot_positions.push_back(Point2(100, 100)); */
40-
/* robot_positions.push_back(Point2(100, 10)); */
41-
/* robot_positions.push_back(Point2(10, 70)); */
42-
43-
/* WorldIDF world_idf(params); */
44-
/* params.pNumRobots = robot_positions.size(); */
45-
/* world_idf.AddNormalDistribution(BivariateNormalDistribution(Point2(95, 95), 10)); */
46-
47-
/* CCT::CoverageSystem env (params, world_idf, robot_positions); */
48-
/* torch::Tensor local_maps = torch::empty({2, params.pNumRobots, params.pLocalMapSize, params.pLocalMapSize}); */
49-
/* std::cout << std::setprecision(10) << std::endl; */
50-
/* env.GetAllRobotsLocalMaps(local_maps[0]); */
51-
/* env.GetAllRobotsLocalMaps(local_maps[1]); */
52-
/* print(local_maps[0].sum()); */
53-
/* std::cout << local_maps.sum() << std::endl; */
54-
/* std::cout << "shape: " << local_maps.sizes(); */
55-
/* torch::save(local_maps[0], "test.pt"); */
56-
57-
/* torch::jit::script::Module resizer_model; */
58-
/* resizer_model = torch::jit::load(resizer_model_path); */
59-
60-
/* std::cout << "shape of local_maps: " << local_maps[0][0].sizes() << std::endl; */
61-
62-
/* std::vector<torch::jit::IValue> inputs; */
63-
/* inputs.push_back(local_maps); */
64-
/* torch::Tensor resized_maps = resizer_model.forward(inputs).toTensor(); */
65-
/* std::cout << "resized_maps: " << resized_maps.sizes() << std::endl; */
66-
67-
/* torch::Tensor comm_maps = torch::zeros({params.pNumRobots, 32, 32}); */
68-
/* env.GetAllRobotsCommunicationMaps(comm_maps, 32); */
69-
/* std::cout << "shape of comm_maps: " << comm_maps.sizes() << std::endl; */
70-
71-
/* auto sparse_comm_maps = comm_maps.to_sparse(); */
72-
/* std::cout << sparse_comm_maps << std::endl; */
73-
74-
/* std::cout << torch::nonzero(comm_maps) - 16 << std::endl; */
75-
76-
/* torch::Tensor test_tensor = torch::zeros({5}); */
77-
/* test_tensor[0] = 1; */
78-
/* test_tensor[1] = 2; */
79-
/* test_tensor[2] = 3; */
80-
/* std::cout << test_tensor << std::endl; */
81-
/* auto test_tensor2 = test_tensor.index({Slice(0, 3)}); */
82-
/* std::cout << test_tensor2 << std::endl; */
83-
/* test_tensor2[0] = 10; */
84-
/* std::cout << test_tensor << std::endl; */
85-
86-
87-
CCT::GenerateDataset dataset_generator("/root/CoverageControl_ws/data/pure_coverage/data_params.yaml");;
23+
if (argc != 2) {
24+
std::cout << "Usage: ./data_generation <path_to_yaml>" << std::endl;
25+
std::cout << "YAML file for data generation is required" << std::endl;
26+
return 1;
27+
}
8828

29+
CCT::GenerateDataset dataset_generator(argv[1]);
8930

9031
return 0;
9132

‎cppsrc/tests/torch_map_conversion.cpp‎

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
#include <CoverageControlTorch/type_conversions.h>
1616

17+
using namespace torch::indexing;
1718
using namespace CoverageControl;
1819
using namespace CoverageControlTorch;
1920

@@ -34,16 +35,31 @@ int main(int argc, char** argv) {
3435
env = std::make_unique<CoverageSystem> (params, world_idf, robot_positions);
3536

3637
auto env_map = env->GetWorldIDF();
37-
auto env_map_tensor = EigenToLibTorch(env_map);
38+
auto env_map_tensor = ToTensor(env_map).clone();
3839
std::cout << env_map_tensor.dtype() << std::endl;
3940

41+
std::cout << env_map_tensor.sizes() << std::endl;
4042
print(env_map_tensor[50][50]);
4143

42-
torch::save(env_map_tensor, "data/test/env_map.pt");
44+
std::cout << "Saving env_map_tensor to env_map.pt" << std::endl;
45+
torch::save(env_map_tensor, "env_map.pt");
4346

4447
torch::Tensor env_map_tensor_loaded;
45-
torch::load(env_map_tensor_loaded, "data/test/env_map.pt");
48+
torch::load(env_map_tensor_loaded, "env_map.pt");
4649
print(env_map_tensor_loaded[50][50]);
4750

51+
52+
std::cout << "Torch indexing and referencing test" << std::endl;
53+
torch::Tensor test_tensor = torch::zeros({5});
54+
test_tensor[0] = 1;
55+
test_tensor[1] = 2;
56+
test_tensor[2] = 3;
57+
std::cout << test_tensor << std::endl;
58+
auto test_tensor2 = test_tensor.index({Slice(0, 3)});
59+
std::cout << test_tensor2 << std::endl;
60+
test_tensor2[0] = 10;
61+
std::cout << test_tensor << std::endl;
62+
63+
4864
return 0;
4965
}

‎cppsrc/torch/include/CoverageControlTorch/coverage_system.h‎

Lines changed: 25 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -13,44 +13,50 @@
1313

1414
#include "type_conversions.h"
1515

16+
typedef long int T_idx_t;
1617
using namespace torch::indexing;
18+
using namespace CoverageControl;
1719
namespace CoverageControlTorch {
1820

1921
class CoverageSystem : public CoverageControl::CoverageSystem {
22+
private:
23+
float env_resolution_ = 1;
24+
float comm_range_ = 256;
25+
26+
void init() {
27+
env_resolution_ = (float) params_.pResolution;
28+
comm_range_ = (float) params_.pCommunicationRange;
29+
}
2030

2131
public:
2232

23-
CoverageSystem(Parameters const &params, size_t const num_features, size_t const num_robots) : CoverageControl::CoverageSystem(params, num_features, num_robots) { }
33+
CoverageSystem(Parameters const &params, size_t const num_features, size_t const num_robots) : CoverageControl::CoverageSystem(params, num_features, num_robots) { init(); }
2434

25-
CoverageSystem(Parameters const &params, WorldIDF const &world_idf, std::string const &pos_file_name) : CoverageControl::CoverageSystem(params, world_idf, pos_file_name) { }
35+
CoverageSystem(Parameters const &params, WorldIDF const &world_idf, std::string const &pos_file_name) : CoverageControl::CoverageSystem(params, world_idf, pos_file_name) { init(); }
2636

27-
CoverageSystem(Parameters const &params, WorldIDF const &world_idf, std::vector <Point2> const &robot_positions) : CoverageControl::CoverageSystem(params, world_idf, robot_positions) { }
37+
CoverageSystem(Parameters const &params, WorldIDF const &world_idf, std::vector <Point2> const &robot_positions) : CoverageControl::CoverageSystem(params, world_idf, robot_positions) { init(); }
2838

29-
CoverageSystem(Parameters const &params, std::vector <BivariateNormalDistribution> const &dists, std::vector <Point2> const &robot_positions) : CoverageControl::CoverageSystem(params, dists, robot_positions) { }
39+
CoverageSystem(Parameters const &params, std::vector <BivariateNormalDistribution> const &dists, std::vector <Point2> const &robot_positions) : CoverageControl::CoverageSystem(params, dists, robot_positions) { init(); }
3040

3141
torch::Tensor GetAllRobotsLocalMaps() {
3242
torch::Tensor maps = torch::zeros({num_robots_, params_.pLocalMapSize, params_.pLocalMapSize});
3343
#pragma omp parallel for
3444
for(size_t i = 0; i < num_robots_; i++) {
35-
maps[i] = EigenToLibTorch(robots_[i].GetRobotLocalMap());
45+
maps[i] = ToTensor(robots_[i].GetRobotLocalMap());
3646
}
3747
return maps;
3848
}
3949

40-
void GetAllRobotsCommunicationMaps(torch::Tensor maps, size_t const &map_size) const {
41-
#pragma omp parallel for
42-
for(int i = 0; i < num_robots_; i++) {
43-
auto robot_neighbors_pos = GetRobotsInCommunication(i);
44-
double comm_scale = (params_.pCommunicationRange * 2.) / map_size;
45-
Point2 map_translation(map_size * comm_scale * params_.pResolution/2., map_size * comm_scale * params_.pResolution/2.);
46-
for(Point2 const& relative_pos:robot_neighbors_pos) {
47-
Point2 map_pos = relative_pos + map_translation;
48-
int pos_idx, pos_idy;
49-
MapUtils::GetClosestGridCoordinate(params_.pResolution * comm_scale, map_pos, pos_idx, pos_idy);
50-
if(pos_idx < map_size and pos_idy < map_size and pos_idx >= 0 and pos_idy >= 0) {
51-
maps.index_put_({i, pos_idx, pos_idy}, 1);
52-
}
53-
}
50+
void GetAllRobotsCommunicationMaps(torch::Tensor maps, size_t const &map_size) {
51+
float f_map_size = (float) map_size;
52+
torch::Tensor robot_positions = ToTensor(GetRobotPositions());
53+
torch::Tensor scaled_relative_pos = torch::round((robot_positions.unsqueeze(0) - robot_positions.unsqueeze(1)) * f_map_size/(comm_range_ * env_resolution_ * 2.) + (f_map_size/2. - env_resolution_/2.)).to(torch::kInt64);
54+
torch::Tensor pairwise_dist_matrices = torch::cdist(robot_positions, robot_positions, 2);
55+
torch::Tensor diagonal_mask = torch::eye(pairwise_dist_matrices.size(0)).to(torch::kBool);
56+
pairwise_dist_matrices.masked_fill_(diagonal_mask, comm_range_ + 1);
57+
for (T_idx_t r_idx = 0; r_idx < num_robots_; ++r_idx) {
58+
torch::Tensor indices = scaled_relative_pos.index({r_idx, pairwise_dist_matrices[r_idx] < (comm_range_ - 0.001), Slice()});
59+
maps.index_put_({r_idx, indices.index({Slice(),0}), indices.index({Slice(), 1})}, 1);
5460
}
5561
}
5662
};

0 commit comments

Comments
 (0)