rotation and other lvl_gen changes

Author: kevin-joseph-ai

src/knn.hpp

@@ -20,7 +20,7 @@ void fillZeros(gpudrive::MapObservation *begin,
                gpudrive::MapObservation *beyond) {
   while (begin < beyond) {
     *begin++ =
-        gpudrive::MapObservation{.position = {0, 0},
+        gpudrive::MapObservation{.position = {0, 0, 0},
                                  .scale = madrona::math::Diag3x3{0, 0, 0},
                                  .heading = 0.f,
                                  .type = (float)gpudrive::EntityType::None};
@@ -30,16 +30,17 @@ void fillZeros(gpudrive::MapObservation *begin,
 gpudrive::MapObservation
 relativeObservation(const gpudrive::MapObservation &absoluteObservation,
                     const madrona::base::Rotation &referenceRotation,
-                    const madrona::math::Vector2 &referencePosition) {
+                    const madrona::math::Vector3 &referencePosition) {
   auto relativePosition =
-      madrona::math::Vector2{.x = absoluteObservation.position.x,
-                             .y = absoluteObservation.position.y} -
+      madrona::math::Vector3{.x = absoluteObservation.position.x,
+                             .y = absoluteObservation.position.y,
+                             .z = absoluteObservation.position.z} -
       referencePosition;
 
   return gpudrive::MapObservation{
       .position = referenceRotation.inv()
-                      .rotateVec({relativePosition.x, relativePosition.y, 0})
-                      .xy(),
+                      .rotateVec({relativePosition.x, relativePosition.y, relativePosition.z})
+                      .xyz(),
       .scale = absoluteObservation.scale,
       .heading =  gpudrive::utils::quatToYaw(referenceRotation.inv() * madrona::math::Quat::angleAxis(absoluteObservation.heading,madrona::math::up)),
       .type = absoluteObservation.type};
@@ -50,7 +51,7 @@ bool isObservationsValid(gpudrive::Engine &ctx,
                          gpudrive::MapObservation *observations,
                          madrona::CountT K,
                          const madrona::base::Rotation &referenceRotation,
-                         const madrona::math::Vector2 &referencePosition) {
+                         const madrona::math::Vector3 &referencePosition) {
 #ifdef MADRONA_GPU_MODE
   return true;
 #else
@@ -102,7 +103,7 @@ namespace gpudrive {
 
 template <madrona::CountT K>
 void selectKNearestRoadEntities(Engine &ctx, const Rotation &referenceRotation,
-                                const madrona::math::Vector2 &referencePosition,
+                                const madrona::math::Vector3 &referencePosition,
                                 gpudrive::MapObservation *heap) {
   const Entity *roads = ctx.data().roads;
   const auto roadCount = ctx.data().numRoads;

src/level_gen.cpp

@@ -163,7 +163,23 @@ static Entity makeRoadEdge(Engine &ctx, const MapRoad &roadInit, CountT j) {
     ctx.get<RoadInterfaceEntity>(road_edge).e = ctx.makeEntity<RoadInterface>();
 
     auto pos = Vector3{.x = (start.x + end.x)/2, .y = (start.y + end.y)/2, .z = (start.z + end.z)/2};
-    auto rot = Quat::angleAxis(atan2(end.y - start.y, end.x - start.x), madrona::math::up);
+
+    //Rotation calculation
+    float dx = end.x - start.x;
+    float dy = end.y - start.y;
+    float dz = end.z - start.z;
+    // XY rotation (yaw)
+    auto yawRot = Quat::angleAxis(atan2(dy, dx), madrona::math::up);
+
+    // Z rotation (pitch)
+    float horizontalDist = sqrt(dx*dx + dy*dy);
+    float pitchAngle = atan2(dz, horizontalDist);
+    Vector3 pitchAxis = Vector3::cross(Vector3{dx/horizontalDist, dy/horizontalDist, 0}, madrona::math::up);
+    auto pitchRot = Quat::angleAxis(pitchAngle, pitchAxis);
+
+    // Combined rotation
+    auto rot = yawRot * pitchRot;
+
     auto scale = Diag3x3{.d0 = start.distance(end)/2, .d1 = 0.1, .d2 = 0.1};
     setRoadEntitiesProps(ctx, road_edge, pos, rot, scale, roadInit.type, ObjectID{(int32_t)SimObject::Cube}, ResponseType::Static, roadInit.id, roadInit.mapType);
     registerRigidBodyEntity(ctx, road_edge, SimObject::Cube);
@@ -367,7 +383,7 @@ void createPersistentEntities(Engine &ctx) {
     ctx.singleton<ResetMap>().reset = 0;
 
     auto& means = ctx.singleton<WorldMeans>().mean;
-    means = {map.mean.x, map.mean.y, 0}; // TODO: Add z to the map
+    means = {map.mean.x, map.mean.y, map.mean.z};
 
     CountT agentIdx = 0;
     for (CountT agentCtr = 0; agentCtr < map.numObjects && agentIdx < consts::kMaxAgentCount; ++agentCtr) {

src/level_gen.hpp

@@ -56,7 +56,7 @@ void destroyWorld(Engine &ctx);
         ctx.get<ResponseType>(road) = responseType;
         ctx.get<RoadMapId>(road).id = roadIdx;
         ctx.get<MapType>(road) = mapType;
-        ctx.get<MapObservation>(ctx.get<RoadInterfaceEntity>(road).e) = MapObservation{.position = pos.xy(),
+        ctx.get<MapObservation>(ctx.get<RoadInterfaceEntity>(road).e) = MapObservation{.position = pos,
                                                                                        .scale = scale,
                                                                                        .heading = utils::quatToYaw(rot),
                                                                                        .type = (float)type,