diff --git a/.gitmodules b/.gitmodules index 2f7c1e26c..f68bb4f2a 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,3 +1,6 @@ [submodule "external/madrona"] path = external/madrona url = https://github.com/shacklettbp/madrona.git +[submodule "external/json"] + path = external/json + url = https://github.com/nlohmann/json.git diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 13677b6bd..5eec975fd 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -33,7 +33,7 @@ target_link_libraries(gpudrive_mgr PUBLIC madrona_python_utils PRIVATE - madrona_json + nlohmann_json::nlohmann_json gpudrive_cpu_impl madrona_mw_cpu madrona_common diff --git a/src/MapReader.cpp b/src/MapReader.cpp index c2fca014f..8c933edd9 100644 --- a/src/MapReader.cpp +++ b/src/MapReader.cpp @@ -36,13 +36,8 @@ madrona_gpudrive::Map *copyToArrayOnHostOrDevice(const madrona_gpudrive::Map *in namespace madrona_gpudrive { -simdjson::dom::parser& MapReader::getParser() { - static simdjson::dom::parser parser; - return parser; -} - -MapReader::MapReader(const std::string &pathToFile){ - in_ = pathToFile; +MapReader::MapReader(const std::string &pathToFile) : in_(pathToFile) { + assert(in_.is_open()); map_ = new madrona_gpudrive::Map(); } @@ -51,13 +46,10 @@ MapReader::~MapReader() { } void MapReader::doParse(float polylineReductionThreshold) { - // Parse with simdjson - auto& parser = getParser(); - simdjson::dom::element doc; - auto error = parser.load(in_).get(doc); + nlohmann::json rawJson; + in_ >> rawJson; - // Parse the document into the map - from_json(doc, *map_, polylineReductionThreshold); + from_json(rawJson, *map_, polylineReductionThreshold); } madrona_gpudrive::Map* MapReader::parseAndWriteOut(const std::string &path, @@ -68,4 +60,4 @@ madrona_gpudrive::Map* MapReader::parseAndWriteOut(const std::string &path, return copyToArrayOnHostOrDevice(reader.map_, executionMode); } -} // namespace madrona_gpudrive \ No newline at end of file +} // namespace madrona_gpudrive diff --git a/src/MapReader.hpp b/src/MapReader.hpp index aac57e458..e42a3f6c6 100644 --- a/src/MapReader.hpp +++ b/src/MapReader.hpp @@ -2,7 +2,7 @@ #include #include -#include + #include "init.hpp" namespace madrona_gpudrive { @@ -14,13 +14,12 @@ class MapReader { static madrona_gpudrive::Map* parseAndWriteOut(const std::string &path, madrona::ExecMode executionMode, float polylineReductionThreshold); private: - static simdjson::dom::parser& getParser(); MapReader(const std::string &pathToFile); ~MapReader(); void doParse(float polylineReductionThreshold); - std::string in_; + std::ifstream in_; madrona_gpudrive::Map *map_; }; -} // namespace madrona_gpudrive \ No newline at end of file +} // namespace madrona_gpudrive diff --git a/src/json_serialization.hpp b/src/json_serialization.hpp index e861d89c2..43cffb837 100644 --- a/src/json_serialization.hpp +++ b/src/json_serialization.hpp @@ -4,242 +4,47 @@ #include "types.hpp" #include "consts.hpp" #include -#include +#include #include namespace madrona_gpudrive { - - #define CHECK_JSON_ERROR(err) \ - if (err) { \ - std::cerr << "JSON error: " << simdjson::error_message(err) << std::endl; \ - abort(); \ - } - - template - T getValueOrDefault(const simdjson::dom::element &obj, - std::string_view key, - T defaultValue) { - simdjson::dom::object json_obj; - auto error = obj.get_object().get(json_obj); - if (error) { - printf("Warning: Failed to get object: %s\n", simdjson::error_message(error)); - abort(); - } - - simdjson::dom::element result; - error = json_obj.at_key(key).get(result); - - if (error == simdjson::NO_SUCH_FIELD) { - printf("Warning: Key %.*s not found in json\n", - static_cast(key.size()), key.data()); - abort(); - } - CHECK_JSON_ERROR(error); - - // Handle different types appropriately - if constexpr (std::is_same_v) { - double d; - error = result.get_double().get(d); - CHECK_JSON_ERROR(error); - return static_cast(d); - } else if constexpr (std::is_same_v) { - int64_t i; - error = result.get_int64().get(i); - CHECK_JSON_ERROR(error); - return static_cast(i); - } else if constexpr (std::is_same_v) { - uint64_t u; - error = result.get_uint64().get(u); - CHECK_JSON_ERROR(error); - return static_cast(u); - } else if constexpr (std::is_same_v) { - uint64_t u; - error = result.get_uint64().get(u); - CHECK_JSON_ERROR(error); - return static_cast(u); - } else if constexpr (std::is_same_v) { - bool b; - error = result.get_bool().get(b); - CHECK_JSON_ERROR(error); - return b; - } else if constexpr (std::is_same_v) { - std::string_view s; - error = result.get_string().get(s); - CHECK_JSON_ERROR(error); - return s; - } else if constexpr (std::is_same_v) { - simdjson::dom::array result_arr; - error = result.get_array().get(result_arr); - CHECK_JSON_ERROR(error); - return result_arr; - } else if constexpr (std::is_same_v) { - return result; - } else { - static_assert(!sizeof(T), "Unsupported type for getValueOrDefault"); - } - } - - template - T getValueOrDefault(const simdjson::dom::array &arr, - size_t idx, - T defaultValue) { - if (idx >= arr.size()) { - printf("Warning: Index %zu not found in array\n", idx); - return defaultValue; - } - - simdjson::dom::element result; - auto error = arr.at(idx).get(result); - CHECK_JSON_ERROR(error); - - // Handle different types appropriately using the same pattern as above - if constexpr (std::is_same_v) { - double d; - error = result.get_double().get(d); - CHECK_JSON_ERROR(error); - return static_cast(d); - } else if constexpr (std::is_same_v) { - return result; - } else if constexpr (std::is_same_v) { - int64_t i; - error = result.get_int64().get(i); - CHECK_JSON_ERROR(error); - return static_cast(i); - } else if constexpr (std::is_same_v) { - uint64_t u; - error = result.get_uint64().get(u); - CHECK_JSON_ERROR(error); - return static_cast(u); - } else if constexpr (std::is_same_v) { - bool b; - error = result.get_bool().get(b); - CHECK_JSON_ERROR(error); - return b; - } else if constexpr (std::is_same_v) { - std::string_view s; - error = result.get_string().get(s); - CHECK_JSON_ERROR(error); - return s; - } else if constexpr (std::is_same_v) { - simdjson::dom::array result_arr; - error = result.get_array().get(result_arr); - CHECK_JSON_ERROR(error); - return result_arr; - } else { - static_assert(!sizeof(T), "Unsupported type for array getValueOrDefault"); - } - } - - void from_dom(const simdjson::dom::element &j, TrafficLightState &tl_state) + void from_json(const nlohmann::json &j, MapVector2 &p) { - // Set number of states to the size of the state array - // simdjson::dom::array states = getValueOrDefault(j, "state", {}); - // simdjson::dom::array x = getValueOrDefault(j, "x", {}); - // simdjson::dom::array y = getValueOrDefault(j, "y", {}); - // simdjson::dom::array z = getValueOrDefault(j, "z", {}); - // simdjson::dom::array timeIndex = getValueOrDefault(j, "time_index", {}); - // simdjson::dom::array laneId = getValueOrDefault(j, "lane_id", {}); - // size_t numStates = std::max(states.size(), static_cast(consts::kTrajectoryLength)); - // tl_state.numStates = numStates; - // static const std::unordered_map state_map = { - // {"unknown", TLState::Unknown}, - // {"stop", TLState::Stop}, - // {"caution", TLState::Caution}, - // {"go", TLState::Go} - // }; - - // Process each timestep - // for (size_t t = 0; t < numStates; t++) { - // // Get the state string and convert to enum - // if (t < states.size()) { - // std::string_view state_str = getValueOrDefault(states, t, "unknown"); - // auto it = state_map.find(state_str); - // TLState enum_state = (it != state_map.end()) ? it->second : TLState::Unknown; // ADD THIS LINE - // tl_state.state[t] = static_cast(enum_state); // Cast enum to float - // } else { - // tl_state.state[t] = static_cast(TLState::Unknown); // Cast enum to float - // } - - // // Get the x,y,z positions in a more interpretable fashion - // if(x.size() > 0) - // { - // tl_state.x[t] = getValueOrDefault(x, 0, 0.0f); - // } - // else - // { - // tl_state.x[t] = -1000.0f; - // } - // if(y.size() > 0) - // { - // tl_state.y[t] = getValueOrDefault(y, 0, 0.0f); - // } - // else - // { - // tl_state.y[t] = -1000.0f; - // } - // if(z.size() > 0) - // { - // tl_state.z[t] = getValueOrDefault(z, 0, 0.0f); - // } - // else - // { - // tl_state.z[t] = -1000.0f; - // } - - // // Get time index and lane id - // if (t < timeIndex.size()) { - // tl_state.timeIndex[t] = getValueOrDefault(timeIndex, t, 0.0f); - // } else { - // tl_state.timeIndex[t] = -1; - // } - // } - - // if (laneId.size() > 0) { - // tl_state.laneId = getValueOrDefault(laneId, 0, 0); - // } - // // Fill any remaining timesteps with default values - // for (size_t t = numStates; t < consts::kTrajectoryLength; t++) { - // tl_state.state[t] = static_cast(TLState::Unknown); - // tl_state.timeIndex[t] = -1; - // } - } - - void from_dom(const simdjson::dom::element &j, MapVector2 &p) { - p.x = getValueOrDefault(j, "x", 0.0f); // Provide a default value of 0.0f if missing - p.y = getValueOrDefault(j, "y", 0.0f); + p.x = j.at("x").get(); + p.y = j.at("y").get(); } - void from_dom(const simdjson::dom::element &j, MapObject &obj) + void from_json(const nlohmann::json &j, MapObject &obj) { obj.mean = {0,0}; uint32_t i = 0; - simdjson::dom::array positions = getValueOrDefault(j, "position", {}); - for (const auto &pos : positions) + for (const auto &pos : j.at("position")) { - if (i >= MAX_POSITIONS) break; // Avoid overflow - from_dom(pos, obj.position[i]); - obj.mean.x += (obj.position[i].x - obj.mean.x)/(i+1); - obj.mean.y += (obj.position[i].y - obj.mean.y)/(i+1); - - ++i; + if (i < MAX_POSITIONS) + { + from_json(pos, obj.position[i]); + obj.mean.x += (obj.position[i].x - obj.mean.x)/(i+1); + obj.mean.y += (obj.position[i].y - obj.mean.y)/(i+1); + ++i; + } + else + { + break; // Avoid overflow + } } obj.numPositions = i; - // Replace with default values (here zero) if the key is missing. - obj.vehicle_size.width = getValueOrDefault(j, "width", 0.0f); - obj.vehicle_size.length = getValueOrDefault(j, "length", 0.0f); - obj.vehicle_size.height = getValueOrDefault(j, "height", 0.0f); - obj.id = getValueOrDefault(j, "id", 0); + j.at("width").get_to(obj.vehicle_size.width); + j.at("length").get_to(obj.vehicle_size.length); + j.at("height").get_to(obj.vehicle_size.height); + j.at("id").get_to(obj.id); i = 0; - simdjson::dom::array headings = getValueOrDefault(j, "heading", {}); - for (const auto &h : headings) + for (const auto &h : j.at("heading")) { if (i < MAX_POSITIONS) { - double heading; - CHECK_JSON_ERROR(h.get(heading)); - obj.heading[i] = static_cast(heading); + h.get_to(obj.heading[i]); ++i; } else @@ -250,31 +55,38 @@ namespace madrona_gpudrive obj.numHeadings = i; i = 0; - simdjson::dom::array velocities = getValueOrDefault(j, "velocity", {}); - for (const auto &v : velocities) + for (const auto &v : j.at("velocity")) { - if (i>= MAX_POSITIONS) break; // Avoid overflow - from_dom(v, obj.velocity[i]); - ++i; + if (i < MAX_POSITIONS) + { + from_json(v, obj.velocity[i]); + ++i; + } + else + { + break; // Avoid overflow + } } obj.numVelocities = i; i = 0; - simdjson::dom::array valid = getValueOrDefault(j, "valid", {}); - for (const auto &v : valid) + for (const auto &v : j.at("valid")) { - if (i >= MAX_POSITIONS) break; // Avoid overflow - bool valid_val; - CHECK_JSON_ERROR(v.get(valid_val)); - obj.valid[i] = valid_val; - ++i; + if (i < MAX_POSITIONS) + { + v.get_to(obj.valid[i]); + ++i; + } + else + { + break; // Avoid overflow + } } obj.numValid = i; - simdjson::dom::element goalPosition = getValueOrDefault(j, "goalPosition", {}); - from_dom(goalPosition, obj.goalPosition); - std::string_view type = getValueOrDefault(j, "type", ""); + from_json(j.at("goalPosition"), obj.goalPosition); + std::string type = j.at("type"); if(type == "vehicle") obj.type = EntityType::Vehicle; else if(type == "pedestrian") @@ -284,7 +96,10 @@ namespace madrona_gpudrive else obj.type = EntityType::None; - obj.markAsExpert = getValueOrDefault(j, "mark_as_expert", false); + std::string markAsExpertKey = "mark_as_expert"; + if (j.contains(markAsExpertKey)) { + from_json(j.at("mark_as_expert"), obj.markAsExpert); + } // Initialize metadata fields to 0 obj.metadata.isSdc = 0; @@ -293,10 +108,10 @@ namespace madrona_gpudrive obj.metadata.difficulty = 0; } - void from_dom(const simdjson::dom::element &j, MapRoad &road, float polylineReductionThreshold = 0.0) + void from_json(const nlohmann::json &j, MapRoad &road, float polylineReductionThreshold = 0.0) { road.mean = {0,0}; - std::string_view type = getValueOrDefault(j, "type", ""); + std::string type = j.at("type"); if(type == "road_edge") road.type = EntityType::RoadEdge; else if(type == "road_line") @@ -312,25 +127,18 @@ namespace madrona_gpudrive else road.type = EntityType::None; + std::vector geometry_points_; - simdjson::dom::array geometry = getValueOrDefault(j, "geometry", {}); - for(const auto &point: geometry) + for(const auto &point: j.at("geometry")) { MapVector2 p; - from_dom(point, p); + from_json(point, p); geometry_points_.push_back(p); } - const int64_t num_segments = geometry.size() - 1; + const int64_t num_segments = j["geometry"].size() - 1; const int64_t sample_every_n_ = 1; const int64_t num_sampled_points = (num_segments + sample_every_n_ - 1) / sample_every_n_ + 1; - - const int64_t geometry_size = static_cast(geometry_points_.size()); - if (geometry_size == 0) { - road.numPoints = 0; - return; - } - if (num_segments >= 10 && (road.type == EntityType::RoadLane || road.type == EntityType::RoadEdge || road.type == EntityType::RoadLine)) { std::vector skip(num_sampled_points, false); // This list tracks the points that are skipped @@ -356,19 +164,9 @@ namespace madrona_gpudrive } if (k_2 >= num_sampled_points) break; - - // Add bounds checking for geometry_points_ access - int64_t idx1 = k * sample_every_n_; - int64_t idx2 = k_1 * sample_every_n_; - int64_t idx3 = k_2 * sample_every_n_; - - if (idx1 >= geometry_size || idx2 >= geometry_size || idx3 >= geometry_size) { - break; - } - - auto point1 = geometry_points_[idx1]; - auto point2 = geometry_points_[idx2]; - auto point3 = geometry_points_[idx3]; + auto point1 = geometry_points_[k * sample_every_n_]; + auto point2 = geometry_points_[k_1 * sample_every_n_]; + auto point3 = geometry_points_[k_2 * sample_every_n_]; float_t area = 0.5 * std::abs((point1.x - point3.x) * (point2.y - point1.y) - (point1.x - point2.x) * (point3.y - point1.y)); if (area < polylineReductionThreshold) { // If the area is less than the threshold, then we skip the middle point @@ -385,54 +183,59 @@ namespace madrona_gpudrive // Create the road lines k = 0; - if (num_sampled_points > 0) { - skip[0] = false; - } - if (num_sampled_points > 1) { - skip[num_sampled_points - 1] = false; - } + skip[0] = false; + skip[num_sampled_points - 1] = false; std::vector new_geometry_points; // This list stores the points that are not skipped while (k < num_sampled_points) { - int64_t idx = k * sample_every_n_; - if (idx < geometry_size && !skip[k]) + if (!skip[k]) { - new_geometry_points.push_back(geometry_points_[idx]); // Add the point to the list if it is not skipped + new_geometry_points.push_back(geometry_points_[k * sample_every_n_]); // Add the point to the list if it is not skipped } k++; } for (size_t i = 0; i < new_geometry_points.size(); i++) { - if(i >= MAX_GEOMETRY) + if(i==MAX_GEOMETRY) break; road.geometry[i] = new_geometry_points[i]; // Create the road lines } - road.numPoints = std::min(new_geometry_points.size(), static_cast(MAX_GEOMETRY)); + road.numPoints = new_geometry_points.size(); } else { for (int64_t i = 0; i < num_sampled_points ; ++i) { - if(i >= MAX_GEOMETRY) + if(i==MAX_GEOMETRY) break; - int64_t idx = i * sample_every_n_; - if (idx >= geometry_size) { - break; - } - road.geometry[i] = geometry_points_[idx]; + road.geometry[i] = geometry_points_[i * sample_every_n_]; } - road.numPoints = std::min(static_cast(num_sampled_points), static_cast(MAX_GEOMETRY)); + road.numPoints = num_sampled_points; } - road.id = getValueOrDefault(j, "id", -1); - - int32_t mapElementId = getValueOrDefault(j, "map_element_id", static_cast(MapType::UNKNOWN)); - if(mapElementId == 4 or mapElementId >= static_cast(MapType::NUM_TYPES) or mapElementId < -1) - road.mapType = MapType::UNKNOWN; + if (j.contains("id")) { + road.id = j.at("id").get(); + } + + if (j.contains("map_element_id")) + { + auto mapElementId = j.at("map_element_id").get(); + + if(mapElementId == 4 or mapElementId >= static_cast(MapType::NUM_TYPES) or mapElementId < -1) + { + road.mapType = MapType::UNKNOWN; + } + else + { + road.mapType = static_cast(mapElementId); + } + } else - road.mapType = static_cast(mapElementId); + { + road.mapType = MapType::UNKNOWN; + } - for (uint32_t i = 0; i < road.numPoints; i++) + for (int i = 0; i < road.numPoints; i++) { road.mean.x += (road.geometry[i].x - road.mean.x)/(i+1); road.mean.y += (road.geometry[i].y - road.mean.y)/(i+1); @@ -440,38 +243,32 @@ namespace madrona_gpudrive } - std::pair calc_mean(const simdjson::dom::element &j) + std::pair calc_mean(const nlohmann::json &j) { std::pair mean = {0, 0}; int64_t numEntities = 0; - simdjson::dom::array objects = getValueOrDefault(j, "objects", {}); - for (const auto &obj : objects) + for (const auto &obj : j["objects"]) { - simdjson::dom::array position = getValueOrDefault(obj, "position", {}); - simdjson::dom::array valids = getValueOrDefault(obj, "valid", {}); int i = 0; - for (const auto &pos : position) { - bool valid_val; - CHECK_JSON_ERROR(valids.at(i).get(valid_val)); - if (valid_val) { - numEntities++; - float newX = getValueOrDefault(pos, "x", 0.0f); - float newY = getValueOrDefault(pos, "y", 0.0f); - mean.first += (newX - mean.first) / numEntities; - mean.second += (newY - mean.second) / numEntities; - } - ++i; + for (const auto &pos : obj["position"]) + { + if(obj["valid"][i++] == false) + continue; + numEntities++; + float newX = pos["x"]; + float newY = pos["y"]; + // Update mean incrementally + mean.first += (newX - mean.first) / numEntities; + mean.second += (newY - mean.second) / numEntities; } } - simdjson::dom::array roads = getValueOrDefault(j, "roads", {}); - for (const auto &obj : roads) + for (const auto &obj : j["roads"]) { - simdjson::dom::array geometry = getValueOrDefault(obj, "geometry", {}); - for (const auto &point : geometry) + for (const auto &point : obj["geometry"]) { numEntities++; - float newX = getValueOrDefault(point, "x", 0.0f); - float newY = getValueOrDefault(point, "y", 0.0f); + float newX = point["x"]; + float newY = point["y"]; // Update mean incrementally mean.first += (newX - mean.first) / numEntities; @@ -481,21 +278,20 @@ namespace madrona_gpudrive return mean; } - void from_json(const simdjson::dom::element &j, Map &map, float polylineReductionThreshold) + void from_json(const nlohmann::json &j, Map &map, float polylineReductionThreshold) { - std::string_view name = getValueOrDefault(j, "name", ""); - std::strncpy(map.mapName, name.data(), sizeof(map.mapName)); + std::string name = j.at("name").get(); + std::strncpy(map.mapName, name.c_str(), sizeof(map.mapName)); - std::string_view scenario_id = getValueOrDefault(j, "scenario_id", ""); - std::strncpy(map.scenarioId, scenario_id.data(), sizeof(map.scenarioId)); + std::string scenario_id = j.at("scenario_id").get(); + std::strncpy(map.scenarioId, scenario_id.c_str(), sizeof(map.scenarioId)); auto mean = calc_mean(j); map.mean = {mean.first, mean.second}; - auto objects = getValueOrDefault(j, "objects", {}); - map.numObjects = std::min(objects.size(), static_cast(MAX_OBJECTS)); + map.numObjects = std::min(j.at("objects").size(), static_cast(MAX_OBJECTS)); - simdjson::dom::element metadata = getValueOrDefault(j, "metadata", {}); - int sdc_index = getValueOrDefault(metadata, "sdc_track_index", -1); + const auto& metadata = j.at("metadata"); + int sdc_index = metadata.at("sdc_track_index").get(); // Create id to object index mapping std::unordered_map idToObjIdx; @@ -506,89 +302,83 @@ namespace madrona_gpudrive std::unordered_set objects_of_interest_ids; // Collect tracks_to_predict indices - simdjson::dom::array tracks_to_predict = getValueOrDefault(metadata, "tracks_to_predict", {}); - for (const auto& track : tracks_to_predict) { - int track_index = getValueOrDefault(track, "track_index", -1); - if (track_index >= 0 && static_cast(track_index) < map.numObjects) { + for (const auto& track : metadata.at("tracks_to_predict")) { + int track_index = track.at("track_index").get(); + if (track_index >= 0 && track_index < j.at("objects").size()) { tracks_to_predict_indices.insert(track_index); } else { - std::cerr << "Warning: Invalid track_index " << track_index << " in scene " << std::string(map.mapName) << std::endl; + std::cerr << "Warning: Invalid track_index " << track_index << " in scene " << j.at("name").get() << std::endl; } } // Collect objects_of_interest IDs - simdjson::dom::array objects_of_interest = getValueOrDefault(metadata, "objects_of_interest", {}); - for (const auto& obj_id : objects_of_interest) { - objects_of_interest_ids.insert(getValueOrDefault(obj_id, "id", -1)); + for (const auto& obj_id : metadata.at("objects_of_interest")) { + objects_of_interest_ids.insert(obj_id.get()); } + // Initialize SDC first if valid - if (sdc_index >= 0 && static_cast(sdc_index) < map.numObjects) { - - simdjson::dom::element sdc_element; - auto error = objects.at(sdc_index).get(sdc_element); - CHECK_JSON_ERROR(error); - from_dom(sdc_element, map.objects[0]); + if (sdc_index >= 0 && sdc_index < j.at("objects").size()) { + j.at("objects")[sdc_index].get_to(map.objects[0]); map.objects[0].metadata.isSdc = 1; // Set additional metadata if needed int sdc_id = map.objects[0].id; if (tracks_to_predict_indices.find(sdc_index) != tracks_to_predict_indices.end()) { - map.objects[0].metadata.isTrackToPredict = 1.0f; + map.objects[0].metadata.isTrackToPredict = 1; // Find and set difficulty - for (const auto& track : tracks_to_predict) { - if (getValueOrDefault(track, "track_index", -1) == sdc_index) { - map.objects[0].metadata.difficulty = getValueOrDefault(track, "difficulty", 0); + for (const auto& track : metadata.at("tracks_to_predict")) { + if (track.at("track_index").get() == sdc_index) { + map.objects[0].metadata.difficulty = track.at("difficulty").get(); break; } } } if (objects_of_interest_ids.find(sdc_id) != objects_of_interest_ids.end()) { - map.objects[0].metadata.isObjectOfInterest = 1.0f; + map.objects[0].metadata.isObjectOfInterest = 1; } - + idToObjIdx[sdc_id] = 0; idx = 1; - + // Remove SDC from sets to avoid double processing tracks_to_predict_indices.erase(sdc_index); objects_of_interest_ids.erase(sdc_id); } // Initialize tracks_to_predict objects (excluding SDC) - for (size_t i = 0; i < map.numObjects && idx < map.numObjects; i++) { - if (i == static_cast(sdc_index)) continue; // Skip SDC as it's already initialized + for (size_t i = 0; i < j.at("objects").size() && idx < map.numObjects; i++) { + if (i == sdc_index) continue; // Skip SDC as it's already initialized if (tracks_to_predict_indices.find(i) != tracks_to_predict_indices.end()) { - from_dom(getValueOrDefault(objects, i, simdjson::dom::element()), map.objects[idx]); + j.at("objects")[i].get_to(map.objects[idx]); map.objects[idx].metadata.isTrackToPredict = 1; // Find and set difficulty - for (const auto& track : tracks_to_predict) { - if (getValueOrDefault(track, "track_index", -1) == static_cast(i)) { - map.objects[idx].metadata.difficulty = getValueOrDefault(track, "difficulty", 0); + for (const auto& track : metadata.at("tracks_to_predict")) { + if (track.at("track_index").get() == static_cast(i)) { + map.objects[idx].metadata.difficulty = track.at("difficulty").get(); break; } } - + // Check if also object of interest if (objects_of_interest_ids.find(map.objects[idx].id) != objects_of_interest_ids.end()) { - map.objects[idx].metadata.isObjectOfInterest = 1.0f; + map.objects[idx].metadata.isObjectOfInterest = 1; objects_of_interest_ids.erase(map.objects[idx].id); } - + idToObjIdx[map.objects[idx].id] = idx; idx++; } } // Initialize objects_of_interest (excluding those already processed) - for (size_t i = 0; i < map.numObjects && idx < map.numObjects; i++) { - if (i == static_cast(sdc_index)) continue; + for (size_t i = 0; i < j.at("objects").size() && idx < map.numObjects; i++) { + if (i == sdc_index) continue; - - int obj_id = getValueOrDefault(getValueOrDefault(objects, i, simdjson::dom::element()), "id", -1); + int obj_id = j.at("objects")[i].at("id").get(); if (objects_of_interest_ids.find(obj_id) != objects_of_interest_ids.end()) { - from_dom(getValueOrDefault(objects, i, simdjson::dom::element()), map.objects[idx]); + j.at("objects")[i].get_to(map.objects[idx]); map.objects[idx].metadata.isObjectOfInterest = 1; idToObjIdx[map.objects[idx].id] = idx; @@ -597,49 +387,29 @@ namespace madrona_gpudrive } // Initialize all remaining objects - for (size_t i = 0; i < map.numObjects && idx < map.numObjects; i++) { - if (i == static_cast(sdc_index)) continue; + for (size_t i = 0; i < j.at("objects").size() && idx < map.numObjects; i++) { + if (i == sdc_index) continue; - int obj_id = getValueOrDefault(getValueOrDefault(objects, i, simdjson::dom::element()), "id", -1); + int obj_id = j.at("objects")[i].at("id").get(); if (idToObjIdx.find(obj_id) == idToObjIdx.end()) { // Check if not already processed - from_dom(getValueOrDefault(objects, i, simdjson::dom::element()), map.objects[idx]); + j.at("objects")[i].get_to(map.objects[idx]); idToObjIdx[map.objects[idx].id] = idx; idx++; } } - + // Process roads - auto roads = getValueOrDefault(j, "roads", {}); - map.numRoads = std::min(roads.size(), static_cast(MAX_ROADS)); + map.numRoads = std::min(j.at("roads").size(), static_cast(MAX_ROADS)); size_t countRoadPoints = 0; idx = 0; - for (const auto &road : roads) { + for (const auto &road : j.at("roads")) { if (idx >= map.numRoads) break; - from_dom(road, map.roads[idx], polylineReductionThreshold); + from_json(road, map.roads[idx], polylineReductionThreshold); size_t roadPoints = map.roads[idx].numPoints; countRoadPoints += (map.roads[idx].type <= EntityType::RoadLane) ? (roadPoints - 1) : 1; ++idx; } map.numRoadSegments = countRoadPoints; - - map.numTrafficLights = 0; - map.hasTrafficLights = false; - auto tl_states_element = getValueOrDefault(j, "tl_states", simdjson::dom::element()); - simdjson::dom::object tl_states; - if (tl_states_element.is_object() && tl_states_element.get_object().get(tl_states) == simdjson::SUCCESS) { - if (std::begin(tl_states) != std::end(tl_states)) { - // non-empty object - map.numTrafficLights = std::min(tl_states.size(), static_cast(consts::kMaxTrafficLightCount)); - map.hasTrafficLights = (map.numTrafficLights > 0); - for (auto [key, value] : tl_states) { - size_t idx = 0; - if (idx >= map.numTrafficLights) - break; - from_dom(value, map.trafficLightStates[idx]); - ++idx; - } - } - } } -} \ No newline at end of file +} diff --git a/tests/test_json_parsing_changes.py b/tests/test_json_parsing_changes.py deleted file mode 100644 index 41d3d8c2c..000000000 --- a/tests/test_json_parsing_changes.py +++ /dev/null @@ -1,67 +0,0 @@ -import os -from pathlib import Path -import mediapy - -# Set working directory to the base directory 'gpudrive' -working_dir = Path.cwd() -while working_dir.name != 'GPUDrive-Fork': - working_dir = working_dir.parent - if working_dir == Path.home(): - raise FileNotFoundError("Base directory 'gpudrive' not found") -os.chdir(working_dir) - -from gpudrive.env.config import EnvConfig, SceneConfig -from gpudrive.env.env_torch import GPUDriveTorchEnv -from gpudrive.visualize.utils import img_from_fig -from gpudrive.env.dataset import SceneDataLoader -import sys -from tqdm import tqdm -import time - -DYNAMICS_MODEL = "delta_local" # "delta_local" / "state" / "classic" -DATA_PATH = sys.argv[1]# Your data path -MAX_NUM_OBJECTS = 64 -NUM_ENVS = 1 -NUM_SCENES = 1000 - -# Configs -env_config = EnvConfig(dynamics_model=DYNAMICS_MODEL) - -# # Make dataloader -data_loader = SceneDataLoader( - root=DATA_PATH, # Path to the dataset - batch_size=NUM_ENVS, # Batch size, you want this to be equal to the number of worlds (envs) so that every world receives a different scene - dataset_size=NUM_SCENES, # Total number of different scenes we want to use - sample_with_replacement=False, - seed=42, - shuffle=True, - file_prefix=sys.argv[2] -) - -# Make environment -env = GPUDriveTorchEnv( - config=env_config, - data_loader=data_loader, - max_cont_agents=MAX_NUM_OBJECTS, # Maximum number of agents to control per scenario - device="cuda", - action_type="continuous" # "continuous" or "discrete" -) - -done_envs = [] -render = False -frames = {f"env_{i}": [] for i in range(NUM_ENVS)} -times = [] -batch_swap_times = [] -# Step through the scene -for i in range(0, data_loader.dataset_size, NUM_ENVS): - - if i < data_loader.dataset_size - NUM_ENVS: - print("Swapping batch: ", i) - batch_swap_start = time.time() - env.swap_data_batch() - batch_swap_elapsed = time.time() - batch_swap_start - batch_swap_times.append(batch_swap_elapsed) - -print(f"Ran {NUM_SCENES} scenes") -print(f"Batch swap total time: {sum(batch_swap_times)}s") -print(f"Batch swap average time: {sum(batch_swap_times) / len(batch_swap_times)}s") \ No newline at end of file