Introduce RoadEntityKNN

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

src/binary_heap.hpp

@@ -0,0 +1,218 @@
+/*
+ *
+ * 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.
+ */
+
+#ifndef __SGI_STL_INTERNAL_HEAP_H
+#define __SGI_STL_INTERNAL_HEAP_H
+
+// Heap-manipulation functions: push_heap, pop_heap, make_heap, sort_heap.
+
+template <class _RandomAccessIterator, class _Distance, class _Tp, 
+          class _Compare>
+
+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 <class _RandomAccessIterator, class _Compare>
+
+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 <class _RandomAccessIterator, class _Compare>
+
+inline void 
+push_heap(_RandomAccessIterator __first, _RandomAccessIterator __last,
+          _Compare __comp)
+{
+  __push_heap_aux(__first, __last, __comp);
+}
+
+template <class _RandomAccessIterator, class _Distance, class _Tp>
+
+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 <class _RandomAccessIterator, class _Tp>
+
+inline void 
+__pop_heap(_RandomAccessIterator __first, _RandomAccessIterator __last,
+           _RandomAccessIterator __result, _Tp __x)
+{
+  *__result = *__first;
+  __adjust_heap(__first, 0, __last - __first, __x);
+}
+
+template <class _RandomAccessIterator>
+
+inline void 
+__pop_heap_aux(_RandomAccessIterator __first, _RandomAccessIterator __last)
+{
+  __pop_heap(__first, __last - 1, __last - 1, *(__last - 1));
+}
+
+template <class _RandomAccessIterator>
+
+inline void pop_heap(_RandomAccessIterator __first, 
+                     _RandomAccessIterator __last)
+{
+  __pop_heap_aux(__first, __last, __VALUE_TYPE(__first));
+}
+
+template <class _RandomAccessIterator, class _Distance,
+          class _Tp, class _Compare>
+
+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 <class _RandomAccessIterator, class _Tp, class _Compare>
+
+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 <class _RandomAccessIterator, class _Compare>
+
+inline void 
+__pop_heap_aux(_RandomAccessIterator __first,
+               _RandomAccessIterator __last, _Compare __comp)
+{
+  __pop_heap(__first, __last - 1, __last - 1, *(__last - 1), __comp);
+}
+
+template <class _RandomAccessIterator, class _Compare>
+
+inline void 
+pop_heap(_RandomAccessIterator __first,
+         _RandomAccessIterator __last, _Compare __comp)
+{
+  __pop_heap_aux(__first, __last, __comp);
+}
+
+template <class _RandomAccessIterator>
+
+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 <class _RandomAccessIterator, class _Compare>
+
+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 <class _RandomAccessIterator, class _Compare>
+
+inline void 
+make_heap(_RandomAccessIterator __first, 
+          _RandomAccessIterator __last, _Compare __comp)
+{
+  __make_heap(__first, __last, __comp);
+}
+
+#endif

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 = 2;
 
 // Various world / entity size parameters
 inline constexpr float worldLength = 40.f;

src/headless.cpp

@@ -61,10 +61,10 @@ int main(int argc, char *argv[])
         .gpuID = 0,
         .numWorlds = (uint32_t)num_worlds,
         .autoReset = false,
-        .jsonPath = "tests/testJsons",
+        .jsonPath = "../maps.150",
         .params = {
             .polylineReductionThreshold = 1.0,
-            .observationRadius = 100.0,
+            .observationRadius = 100000.0,
             .rewardParams = {
                 .rewardType = RewardType::DistanceBased,
                 .distanceToGoalThreshold = 0.5,
@@ -151,7 +151,7 @@ int main(int argc, char *argv[])
             }
         }
         mgr.step();
-        printObs();
+	//        printObs();
     }
 
     auto end = std::chrono::system_clock::now();

src/knn.hpp

@@ -0,0 +1,136 @@
+#include <madrona/math.hpp>
+#include <madrona/types.hpp>
+#include "binary_heap.hpp"
+#include "types.hpp"
+#include <algorithm>
+
+#include <vector>
+
+namespace {
+  struct Comparator {
+    Comparator() {}
+    bool operator()(const gpudrive::MapObservation& lhs, const gpudrive::MapObservation& rhs) const {
+      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<gpudrive::MapObservation> sortedObservations;
+    sortedObservations.reserve(roadCount);
+
+    for (madrona::CountT roadIdx = 0; roadIdx < roadCount; ++roadIdx) {
+      const auto& currentObservation = ctx.get<gpudrive::MapObservation>(ctx.data().roads[roadIdx]);
+      sortedObservations.emplace_back(relativeObservation(currentObservation,
+							  referenceRotation,
+							  referencePosition));
+    }
+
+    Comparator cmp;
+    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
+    return true;
+  }
+}
+
+namespace gpudrive {
+
+template <madrona::CountT K>
+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<gpudrive::MapObservation>(roads[i]),
+				  referenceRotation,
+				  referencePosition);
+  }
+
+  if (roadCount < K) {
+    fillZeros(heap + roadCount, heap + K);
+    return;
+  }
+
+  Comparator cmp;
+  make_heap(heap, heap + K, cmp);
+
+  for (madrona::CountT roadIdx = K; roadIdx < roadCount; ++roadIdx) {
+    auto currentObservation = relativeObservation(ctx.get<gpudrive::MapObservation>(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);
+}
+
+}