From d827b00623259668748367d728593e405cf0d04f Mon Sep 17 00:00:00 2001 From: SamanKazemkhani Date: Thu, 2 May 2024 14:40:40 -0400 Subject: [PATCH] Introduce RoadEntityKNN Add feature flag --- src/CMakeLists.txt | 2 + src/binary_heap.hpp | 192 ++++++++++++++++++++++++++++++++++++++ src/consts.hpp | 5 +- src/headless.cpp | 7 +- src/init.hpp | 7 ++ src/knn.hpp | 140 +++++++++++++++++++++++++++ src/mgr.cpp | 14 +++ src/sim.cpp | 11 ++- src/types.hpp | 7 +- tests/observationTest.cpp | 1 + 10 files changed, 375 insertions(+), 11 deletions(-) create mode 100644 src/binary_heap.hpp create mode 100644 src/knn.hpp diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index b119ce878..cb56fa13d 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -5,6 +5,8 @@ set(SIMULATOR_SRCS level_gen.hpp level_gen.cpp obb.hpp utils.hpp + binary_heap.hpp + knn.hpp ) add_library(gpudrive_cpu_impl STATIC diff --git a/src/binary_heap.hpp b/src/binary_heap.hpp new file mode 100644 index 000000000..05bc3e03f --- /dev/null +++ b/src/binary_heap.hpp @@ -0,0 +1,192 @@ +/* + * + * Copyright (c) 1994 + * Hewlett-Packard Company + * + * Permission to use, copy, modify, distribute and sell this software + * and its documentation for any purpose is hereby granted without fee, + * provided that the above copyright notice appear in all copies and + * that both that copyright notice and this permission notice appear + * in supporting documentation. Hewlett-Packard Company makes no + * representations about the suitability of this software for any + * purpose. It is provided "as is" without express or implied warranty. + * + * Copyright (c) 1997 + * Silicon Graphics Computer Systems, Inc. + * + * Permission to use, copy, modify, distribute and sell this software + * and its documentation for any purpose is hereby granted without fee, + * provided that the above copyright notice appear in all copies and + * that both that copyright notice and this permission notice appear + * in supporting documentation. Silicon Graphics makes no + * representations about the suitability of this software for any + * purpose. It is provided "as is" without express or implied warranty. + */ + +/* NOTE: This is an internal header file, included by other STL headers. + * You should not attempt to use it directly. + */ +#pragma once + +// Heap-manipulation functions: push_heap, pop_heap, make_heap, sort_heap. + +template + +inline void __push_heap(_RandomAccessIterator __first, _Distance __holeIndex, + _Distance __topIndex, _Tp __x, _Compare __comp) { + _Distance __parent = (__holeIndex - 1) / 2; + while (__holeIndex > __topIndex && __comp(*(__first + __parent), __x)) { + *(__first + __holeIndex) = *(__first + __parent); + __holeIndex = __parent; + __parent = (__holeIndex - 1) / 2; + } + *(__first + __holeIndex) = __x; +} + +template + +inline void __push_heap_aux(_RandomAccessIterator __first, + _RandomAccessIterator __last, _Compare __comp) { + int hole_index = (__last - __first) - 1; + int top_index = 0; + __push_heap(__first, hole_index, top_index, *(__last - 1), __comp); +} + +template + +inline void push_heap(_RandomAccessIterator __first, + _RandomAccessIterator __last, _Compare __comp) { + __push_heap_aux(__first, __last, __comp); +} + +template + +void __adjust_heap(_RandomAccessIterator __first, _Distance __holeIndex, + _Distance __len, _Tp __x) { + _Distance __topIndex = __holeIndex; + _Distance __secondChild = 2 * __holeIndex + 2; + while (__secondChild < __len) { + if (*(__first + __secondChild) < *(__first + (__secondChild - 1))) + __secondChild--; + *(__first + __holeIndex) = *(__first + __secondChild); + __holeIndex = __secondChild; + __secondChild = 2 * (__secondChild + 1); + } + if (__secondChild == __len) { + *(__first + __holeIndex) = *(__first + (__secondChild - 1)); + __holeIndex = __secondChild - 1; + } + __push_heap(__first, __holeIndex, __topIndex, __x); +} + +template + +inline void __pop_heap(_RandomAccessIterator __first, + _RandomAccessIterator __last, + _RandomAccessIterator __result, _Tp __x) { + *__result = *__first; + __adjust_heap(__first, 0, __last - __first, __x); +} + +template + +inline void __pop_heap_aux(_RandomAccessIterator __first, + _RandomAccessIterator __last) { + __pop_heap(__first, __last - 1, __last - 1, *(__last - 1)); +} + +template + +inline void pop_heap(_RandomAccessIterator __first, + _RandomAccessIterator __last) { + __pop_heap_aux(__first, __last, __VALUE_TYPE(__first)); +} + +template + +void __adjust_heap(_RandomAccessIterator __first, _Distance __holeIndex, + _Distance __len, _Tp __x, _Compare __comp) { + _Distance __topIndex = __holeIndex; + _Distance __secondChild = 2 * __holeIndex + 2; + while (__secondChild < __len) { + if (__comp(*(__first + __secondChild), *(__first + (__secondChild - 1)))) + __secondChild--; + *(__first + __holeIndex) = *(__first + __secondChild); + __holeIndex = __secondChild; + __secondChild = 2 * (__secondChild + 1); + } + if (__secondChild == __len) { + *(__first + __holeIndex) = *(__first + (__secondChild - 1)); + __holeIndex = __secondChild - 1; + } + __push_heap(__first, __holeIndex, __topIndex, __x, __comp); +} + +template + +inline void +__pop_heap(_RandomAccessIterator __first, _RandomAccessIterator __last, + _RandomAccessIterator __result, _Tp __x, _Compare __comp) { + *__result = *__first; + + int hole_index = 0; + int len = __last - __first; + + __adjust_heap(__first, hole_index, len, __x, __comp); +} + +template + +inline void __pop_heap_aux(_RandomAccessIterator __first, + _RandomAccessIterator __last, _Compare __comp) { + __pop_heap(__first, __last - 1, __last - 1, *(__last - 1), __comp); +} + +template + +inline void pop_heap(_RandomAccessIterator __first, + _RandomAccessIterator __last, _Compare __comp) { + __pop_heap_aux(__first, __last, __comp); +} + +template + +void __make_heap(_RandomAccessIterator __first, _RandomAccessIterator __last) { + if (__last - __first < 2) + return; + int __len = __last - __first; + int __parent = (__len - 2) / 2; + + while (true) { + __adjust_heap(__first, __parent, __len, *(__first + __parent)); + if (__parent == 0) + return; + __parent--; + } +} + +template + +void __make_heap(_RandomAccessIterator __first, _RandomAccessIterator __last, + _Compare __comp) { + if (__last - __first < 2) + return; + int __len = __last - __first; + int __parent = (__len - 2) / 2; + + while (true) { + __adjust_heap(__first, __parent, __len, *(__first + __parent), __comp); + if (__parent == 0) + return; + __parent--; + } +} + +template + +inline void make_heap(_RandomAccessIterator __first, + _RandomAccessIterator __last, _Compare __comp) { + __make_heap(__first, __last, __comp); +} diff --git a/src/consts.hpp b/src/consts.hpp index edabe8c11..80a4b19f6 100644 --- a/src/consts.hpp +++ b/src/consts.hpp @@ -8,8 +8,9 @@ namespace gpudrive { namespace consts { -inline constexpr madrona::CountT kMaxAgentCount = 50; -inline constexpr madrona::CountT kMaxRoadEntityCount = 6000; +inline constexpr madrona::CountT kMaxAgentCount = 300; +inline constexpr madrona::CountT kMaxRoadEntityCount = 3000; +inline constexpr madrona::CountT kMaxAgentMapObservationsCount = 256; // Various world / entity size parameters inline constexpr float worldLength = 40.f; diff --git a/src/headless.cpp b/src/headless.cpp index 2ef36c735..60cd1ea8b 100755 --- a/src/headless.cpp +++ b/src/headless.cpp @@ -61,7 +61,7 @@ int main(int argc, char *argv[]) .gpuID = 0, .numWorlds = (uint32_t)num_worlds, .autoReset = false, - .jsonPath = "tests/testJsons", + .jsonPath = "../maps.16", .params = { .polylineReductionThreshold = 1.0, .observationRadius = 100.0, @@ -71,8 +71,7 @@ int main(int argc, char *argv[]) .distanceToExpertThreshold = 0.5 }, .maxNumControlledVehicles = 0 - }, - .enableBatchRenderer = true + } }); std::random_device rd; @@ -127,7 +126,6 @@ int main(int argc, char *argv[]) printf("Info\n"); info_printer.print(); }; - // printObs(); auto worldToShape = mgr.getShapeTensorFromDeviceMemory(exec_mode, num_worlds); @@ -152,7 +150,6 @@ int main(int argc, char *argv[]) } } mgr.step(); - printObs(); } const auto end = std::chrono::steady_clock::now(); const std::chrono::duration elapsed = end - start; diff --git a/src/init.hpp b/src/init.hpp index 11593f846..043c5e35c 100755 --- a/src/init.hpp +++ b/src/init.hpp @@ -95,6 +95,11 @@ namespace gpudrive Ignore }; + enum class FindRoadObservationsWith { + KNearestEntitiesWithRadiusFiltering, + AllEntitiesWithRadiusFiltering + }; + struct Parameters { float polylineReductionThreshold; @@ -104,6 +109,8 @@ namespace gpudrive CollisionBehaviour collisionBehaviour = CollisionBehaviour::AgentStop; // Default: AgentStop uint32_t maxNumControlledVehicles = 10000; // Arbitrary high number to by default control all vehicles bool IgnoreNonVehicles = false; // Default: false + FindRoadObservationsWith roadObservationAlgorithm{ + FindRoadObservationsWith::KNearestEntitiesWithRadiusFiltering}; }; struct WorldInit diff --git a/src/knn.hpp b/src/knn.hpp new file mode 100644 index 000000000..42584c97c --- /dev/null +++ b/src/knn.hpp @@ -0,0 +1,140 @@ +#pragma once + +#include "binary_heap.hpp" +#include "types.hpp" +#include +#include +#include + +#ifndef MADRONA_GPU_MODE +#include +#endif + +namespace { +bool cmp(const gpudrive::MapObservation &lhs, const gpudrive::MapObservation &rhs) { + return lhs.position.length2() < rhs.position.length2(); +} + +void fillZeros(gpudrive::MapObservation *begin, + gpudrive::MapObservation *beyond) { + while (begin < beyond) { + *begin++ = + gpudrive::MapObservation{.position = {0, 0}, + .scale = madrona::math::Diag3x3{0, 0, 0}, + .heading = 0.f, + .type = (float)gpudrive::EntityType::None}; + } +} + +gpudrive::MapObservation +relativeObservation(const gpudrive::MapObservation &absoluteObservation, + const madrona::base::Rotation &referenceRotation, + const madrona::math::Vector2 &referencePosition) { + auto relativePosition = + madrona::math::Vector2{.x = absoluteObservation.position.x, + .y = absoluteObservation.position.y} - + referencePosition; + + return gpudrive::MapObservation{ + .position = referenceRotation.inv() + .rotateVec({relativePosition.x, relativePosition.y, 0}) + .xy(), + .scale = absoluteObservation.scale, + .heading = absoluteObservation.heading, + .type = absoluteObservation.type}; +} + +bool isObservationsValid(gpudrive::Engine &ctx, + gpudrive::MapObservation *observations, + madrona::CountT K, + const madrona::base::Rotation &referenceRotation, + const madrona::math::Vector2 &referencePosition) { +#ifdef MADRONA_GPU_MODE + return true; +#else + const auto roadCount = ctx.data().numRoads; + + std::vector sortedObservations; + sortedObservations.reserve(roadCount); + + for (madrona::CountT roadIdx = 0; roadIdx < roadCount; ++roadIdx) { + const auto ¤tObservation = + ctx.get(ctx.data().roads[roadIdx]); + sortedObservations.emplace_back(relativeObservation( + currentObservation, referenceRotation, referencePosition)); + } + + std::sort(sortedObservations.begin(), sortedObservations.end(), cmp); + std::sort(observations, observations + K, cmp); + + return std::equal(observations, observations + K, sortedObservations.begin(), + sortedObservations.begin() + K, + [](const gpudrive::MapObservation &lhs, + const gpudrive::MapObservation &rhs) { + return lhs.position.x == rhs.position.x && + lhs.position.y == rhs.position.y; + }); +#endif +} +} // namespace + +namespace gpudrive { + +template +void selectKNearestRoadEntities(Engine &ctx, const Rotation &referenceRotation, + const madrona::math::Vector2 &referencePosition, + gpudrive::MapObservation *heap) { + const Entity *roads = ctx.data().roads; + const auto roadCount = ctx.data().numRoads; + + for (madrona::CountT i = 0; i < std::min(roadCount, K); ++i) { + heap[i] = relativeObservation(ctx.get(roads[i]), + referenceRotation, referencePosition); + } + + if (roadCount < K) { + fillZeros(heap + roadCount, heap + K); + return; + } + + make_heap(heap, heap + K, cmp); + + for (madrona::CountT roadIdx = K; roadIdx < roadCount; ++roadIdx) { + auto currentObservation = + relativeObservation(ctx.get(roads[roadIdx]), + referenceRotation, referencePosition); + + const auto &kthNearestObservation = heap[0]; + bool isCurrentObservationCloser = + cmp(currentObservation, kthNearestObservation); + + if (not isCurrentObservationCloser) { + continue; + } + + pop_heap(heap, heap + K, cmp); + + heap[K - 1] = currentObservation; + push_heap(heap, heap + K, cmp); + } + + assert( + isObservationsValid(ctx, heap, K, referenceRotation, referencePosition)); + + madrona::CountT newBeyond{K}; + { + madrona::CountT idx{0}; + while (idx < newBeyond) { + if (heap[idx].position.length() <= ctx.data().params.observationRadius) { + ++idx; + continue; + } + + heap[idx] = heap[--newBeyond]; + } + } + + fillZeros(heap + newBeyond, heap + K); +} + +} // namespace gpudrive diff --git a/src/mgr.cpp b/src/mgr.cpp index 82807a575..c1f8ff9f6 100755 --- a/src/mgr.cpp +++ b/src/mgr.cpp @@ -437,6 +437,17 @@ static std::vector getMapFiles(const Manager::Config &cfg) return mapFiles; } +bool isRoadObservationAlgorithmValid(FindRoadObservationsWith algo) { + madrona::CountT roadObservationsCount = + sizeof(AgentMapObservations) / sizeof(MapObservation); + + return algo == + FindRoadObservationsWith::KNearestEntitiesWithRadiusFiltering || + (algo == + FindRoadObservationsWith::AllEntitiesWithRadiusFiltering && + roadObservationsCount == consts::kMaxRoadEntityCount); +} + Manager::Impl * Manager::Impl::init( const Manager::Config &mgr_cfg) { @@ -445,6 +456,9 @@ Manager::Impl * Manager::Impl::init( std::vector mapFiles = getMapFiles(mgr_cfg); + assert(isRoadObservationAlgorithmValid( + mgr_cfg.params.roadObservationAlgorithm)); + switch (mgr_cfg.execMode) { case ExecMode::CUDA: { #ifdef MADRONA_CUDA_SUPPORT diff --git a/src/sim.cpp b/src/sim.cpp index 7d22fb477..00790f9e1 100755 --- a/src/sim.cpp +++ b/src/sim.cpp @@ -7,6 +7,7 @@ #include "obb.hpp" #include "sim.hpp" #include "utils.hpp" +#include "knn.hpp" using namespace madrona; using namespace madrona::math; @@ -212,6 +213,14 @@ inline void collectObservationsSystem(Engine &ctx, arrIndex++; } + const auto alg = ctx.data().params.roadObservationAlgorithm; + if (alg == FindRoadObservationsWith::KNearestEntitiesWithRadiusFiltering) { + selectKNearestRoadEntities( + ctx, rot, model.position, map_obs.obs); + return; + } + + assert(alg == FindRoadObservationsWith::AllEntitiesWithRadiusFiltering); arrIndex = 0; CountT roadIdx = 0; while(roadIdx < ctx.data().numRoads) { Entity road = ctx.data().roads[roadIdx++]; @@ -697,7 +706,7 @@ void Sim::setupTasks(TaskGraphManager &taskgraph_mgr, const Config &cfg) auto detectCollisions = builder.addToGraph< ParallelForNode>( - {broadphase_setup_sys}); + {findOverlappingEntities}); // Improve controllability of agents by setting their velocity to 0 // after physics is done. diff --git a/src/types.hpp b/src/types.hpp index 1e5721c28..8140c9232 100644 --- a/src/types.hpp +++ b/src/types.hpp @@ -138,13 +138,14 @@ static_assert(sizeof(PartnerObservations) == sizeof(float) * (consts::kMaxAgentCount - 1) * PartnerObservationExportSize); struct AgentMapObservations { - MapObservation obs[consts::kMaxRoadEntityCount]; + MapObservation obs[consts::kMaxAgentMapObservationsCount]; }; const size_t AgentMapObservationExportSize = 7; -static_assert(sizeof(AgentMapObservations) == sizeof(float) * - consts::kMaxRoadEntityCount * AgentMapObservationExportSize); +static_assert(sizeof(AgentMapObservations) == + sizeof(float) * consts::kMaxAgentMapObservationsCount * + AgentMapObservationExportSize); struct LidarSample { float depth; diff --git a/tests/observationTest.cpp b/tests/observationTest.cpp index c1fe89c25..70f0fd8b6 100755 --- a/tests/observationTest.cpp +++ b/tests/observationTest.cpp @@ -32,6 +32,7 @@ class ObservationsTest : public ::testing::Test { .polylineReductionThreshold = 0.0, .observationRadius = 100.0, .collisionBehaviour = gpudrive::CollisionBehaviour::Ignore, + .roadObservationAlgorithm = gpudrive::FindRoadObservationsWith::AllEntitiesWithRadiusFiltering } });