diff --git a/src/init.hpp b/src/init.hpp index 7c4581687..e87c8ca75 100755 --- a/src/init.hpp +++ b/src/init.hpp @@ -13,40 +13,41 @@ namespace gpudrive // Cannot use Madrona::math::Vector2 because it is not a POD type. // Getting all zeros if using any madrona types. - struct MapVector2 + struct MapVector3 { float x; float y; + float z; }; struct MapObject { - MapVector2 position[MAX_POSITIONS]; + MapVector3 position[MAX_POSITIONS]; float width; float length; float heading[MAX_POSITIONS]; - MapVector2 velocity[MAX_POSITIONS]; + MapVector3 velocity[MAX_POSITIONS]; bool valid[MAX_POSITIONS]; - MapVector2 goalPosition; + MapVector3 goalPosition; EntityType type; uint32_t numPositions; uint32_t numHeadings; uint32_t numVelocities; uint32_t numValid; - MapVector2 mean; + MapVector3 mean; bool markAsExpert{false}; }; struct MapRoad { // std::array geometry; - MapVector2 geometry[MAX_GEOMETRY]; + MapVector3 geometry[MAX_GEOMETRY]; uint32_t id; MapType mapType; EntityType type; uint32_t numPoints; - MapVector2 mean; + MapVector3 mean; }; struct Map @@ -57,7 +58,7 @@ namespace gpudrive uint32_t numObjects; uint32_t numRoads; uint32_t numRoadSegments; - MapVector2 mean; + MapVector3 mean; // Constructor Map() = default; @@ -112,7 +113,7 @@ namespace gpudrive uint32_t maxNumControlledAgents = 10000; // Arbitrary high number to by default control all vehicles bool IgnoreNonVehicles = false; // Default: false FindRoadObservationsWith roadObservationAlgorithm{ - FindRoadObservationsWith::KNearestEntitiesWithRadiusFiltering}; + FindRoadObservationsWith::AllEntitiesWithRadiusFiltering}; bool initOnlyValidAgentsAtFirstStep = true; // Default: true bool isStaticAgentControlled = false; // Default: false bool enableLidar = false; diff --git a/src/json_serialization.hpp b/src/json_serialization.hpp index 40e8762cc..706d346dd 100644 --- a/src/json_serialization.hpp +++ b/src/json_serialization.hpp @@ -7,15 +7,21 @@ namespace gpudrive { - void from_json(const nlohmann::json &j, MapVector2 &p) + void from_json(const nlohmann::json &j, MapVector3 &p) { p.x = j.at("x").get(); p.y = j.at("y").get(); + if (j.contains("z")) { + p.z = j.at("z").get(); + } + else { + p.z = 0.0; + } } void from_json(const nlohmann::json &j, MapObject &obj) { - obj.mean = {0,0}; + obj.mean = {0,0,0}; uint32_t i = 0; for (const auto &pos : j.at("position")) { @@ -24,6 +30,7 @@ namespace gpudrive 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); + obj.mean.z += (obj.position[i].z - obj.mean.z)/(i+1); ++i; } else @@ -100,7 +107,7 @@ namespace gpudrive void from_json(const nlohmann::json &j, MapRoad &road, float polylineReductionThreshold = 0.0) { - road.mean = {0,0}; + road.mean = {0,0,0}; std::string type = j.at("type"); if(type == "road_edge") road.type = EntityType::RoadEdge; @@ -118,10 +125,10 @@ namespace gpudrive road.type = EntityType::None; - std::vector geometry_points_; + std::vector geometry_points_; for(const auto &point: j.at("geometry")) { - MapVector2 p; + MapVector3 p; from_json(point, p); geometry_points_.push_back(p); } @@ -154,10 +161,24 @@ namespace gpudrive } if (k_2 >= num_sampled_points) break; + //Using 3D space area calculation 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)); + // Vector 1: point1 -> point2 + float vx1 = point2.x - point1.x; + float vy1 = point2.y - point1.y; + float vz1 = point2.z - point1.z; + // Vector 2: point1 -> point3 + float vx2 = point3.x - point1.x; + float vy2 = point3.y - point1.y; + float vz2 = point3.z - point1.z; + // Cross product to get the normal vector + float nx = vy1 * vz2 - vz1 * vy2; + float ny = vz1 * vx2 - vx1 * vz2; + float nz = vx1 * vy2 - vy1 * vx2; + // Magnitude of the cross product + float_t area = 0.5 * std::sqrt(nx * nx + ny * ny + nz * nz); if (area < polylineReductionThreshold) { // If the area is less than the threshold, then we skip the middle point skip[k_1] = true; // Mark the middle point as skipped @@ -175,7 +196,7 @@ namespace gpudrive k = 0; skip[0] = false; skip[num_sampled_points - 1] = false; - std::vector new_geometry_points; // This list stores the points that are not skipped + std::vector new_geometry_points; // This list stores the points that are not skipped while (k < num_sampled_points) { if (!skip[k]) @@ -229,14 +250,16 @@ namespace gpudrive { road.mean.x += (road.geometry[i].x - road.mean.x)/(i+1); road.mean.y += (road.geometry[i].y - road.mean.y)/(i+1); + road.mean.z += (road.geometry[i].z - road.mean.z)/(i+1); } } - std::pair calc_mean(const nlohmann::json &j) + std::tuple calc_mean(const nlohmann::json &j) { - std::pair mean = {0, 0}; + std::tuple mean = {0, 0, 0}; int64_t numEntities = 0; + for (const auto &obj : j["objects"]) { int i = 0; @@ -247,11 +270,15 @@ namespace gpudrive numEntities++; float newX = pos["x"]; float newY = pos["y"]; + float newZ = pos["z"]; + // Update mean incrementally - mean.first += (newX - mean.first) / numEntities; - mean.second += (newY - mean.second) / numEntities; + std::get<0>(mean) += (newX - std::get<0>(mean)) / numEntities; + std::get<1>(mean) += (newY - std::get<1>(mean)) / numEntities; + std::get<2>(mean) += (newZ - std::get<2>(mean)) / numEntities; } } + for (const auto &obj : j["roads"]) { for (const auto &point : obj["geometry"]) @@ -259,12 +286,14 @@ namespace gpudrive numEntities++; float newX = point["x"]; float newY = point["y"]; - - // Update mean incrementally - mean.first += (newX - mean.first) / numEntities; - mean.second += (newY - mean.second) / numEntities; + float newZ = point["z"]; + + std::get<0>(mean) += (newX - std::get<0>(mean)) / numEntities; + std::get<1>(mean) += (newY - std::get<1>(mean)) / numEntities; + std::get<2>(mean) += (newZ - std::get<2>(mean)) / numEntities; } } + return mean; } @@ -272,7 +301,8 @@ namespace gpudrive void from_json(const nlohmann::json &j, Map &map, float polylineReductionThreshold) { auto mean = calc_mean(j); - map.mean = {mean.first, mean.second}; + map.mean = {std::get<0>(mean), std::get<1>(mean), std::get<2>(mean)}; // Updated to handle all 3 components + map.numObjects = std::min(j.at("objects").size(), static_cast(MAX_OBJECTS)); size_t idx = 0; for (const auto &obj : j.at("objects")) @@ -281,7 +311,7 @@ namespace gpudrive break; obj.get_to(map.objects[idx++]); } - + map.numRoads = std::min(j.at("roads").size(), static_cast(MAX_ROADS)); size_t countRoadPoints = 0; idx = 0; diff --git a/src/knn.hpp b/src/knn.hpp index c63bc9663..51822ad89 100644 --- a/src/knn.hpp +++ b/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,16 @@ 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}), .scale = absoluteObservation.scale, .heading = gpudrive::utils::quatToYaw(referenceRotation.inv() * madrona::math::Quat::angleAxis(absoluteObservation.heading,madrona::math::up)), .type = absoluteObservation.type}; @@ -50,7 +50,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 +102,7 @@ namespace gpudrive { template 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; diff --git a/src/level_gen.cpp b/src/level_gen.cpp index add4cc4f5..a6ad4932c 100755 --- a/src/level_gen.cpp +++ b/src/level_gen.cpp @@ -33,11 +33,12 @@ static inline void resetAgent(Engine &ctx, Entity agent) { auto agent_iface = ctx.get(agent).e; auto xCoord = ctx.get(agent_iface).positions[0].x; auto yCoord = ctx.get(agent_iface).positions[0].y; + auto zCoord = ctx.get(agent_iface).positions[0].z; auto xVelocity = ctx.get(agent_iface).velocities[0].x; auto yVelocity = ctx.get(agent_iface).velocities[0].y; auto heading = ctx.get(agent_iface).headings[0]; - ctx.get(agent) = Vector3{.x = xCoord, .y = yCoord, .z = 1}; + ctx.get(agent) = Vector3{.x = xCoord, .y = yCoord, .z = zCoord}; ctx.get(agent) = Quat::angleAxis(heading, madrona::math::up); if (ctx.get(agent) == ResponseType::Static) { ctx.get(agent) = Velocity{Vector3::zero(), Vector3::zero()}; @@ -58,7 +59,11 @@ static inline void populateExpertTrajectory(Engine &ctx, const Entity &agent, co auto &trajectory = ctx.get(agent_iface); for(CountT i = 0; i < agentInit.numPositions; i++) { - trajectory.positions[i] = Vector2{.x = agentInit.position[i].x - ctx.singleton().mean.x, .y = agentInit.position[i].y - ctx.singleton().mean.y}; + trajectory.positions[i] = Vector3{ + .x = agentInit.position[i].x - ctx.singleton().mean.x, + .y = agentInit.position[i].y - ctx.singleton().mean.y, + .z = agentInit.position[i].z - ctx.singleton().mean.z + }; trajectory.velocities[i] = Vector2{.x = agentInit.velocity[i].x, .y = agentInit.velocity[i].y}; trajectory.headings[i] = toRadians(agentInit.heading[i]); trajectory.valids[i] = (float)agentInit.valid[i]; @@ -75,7 +80,7 @@ static inline void populateExpertTrajectory(Engine &ctx, const Entity &agent, co } Rotation rot = Quat::angleAxis(trajectory.headings[i], madrona::math::up); - Position pos = Vector3{.x = trajectory.positions[i].x, .y = trajectory.positions[i].y, .z = 1}; + Position pos = Vector3{.x = trajectory.positions[i].x, .y = trajectory.positions[i].y, .z = trajectory.positions[i].z}; Velocity vel = {Vector3{.x = trajectory.velocities[i].x, .y = trajectory.velocities[i].y, .z = 0}, Vector3::zero()}; Rotation targetRot = Quat::angleAxis(trajectory.headings[i+1], madrona::math::up); switch (ctx.data().params.dynamicsModel) { @@ -91,7 +96,7 @@ static inline void populateExpertTrajectory(Engine &ctx, const Entity &agent, co } case DynamicsModel::DeltaLocal: { - Position targetPos = Vector3{.x = trajectory.positions[i+1].x, .y = trajectory.positions[i+1].y, .z = 1}; + Position targetPos = Vector3{.x = trajectory.positions[i+1].x, .y = trajectory.positions[i+1].y, .z = trajectory.positions[i+1].z}; trajectory.inverseActions[i] = inverseDeltaModel(rot, pos, targetRot, targetPos); break; } @@ -126,8 +131,10 @@ static inline Entity createAgent(Engine &ctx, const MapObject &agentInit) { ctx.get(agent) *= consts::vehicleLengthScale; ctx.get(agent) = ObjectID{(int32_t)SimObject::Agent}; ctx.get(agent) = agentInit.type; - ctx.get(agent)= Goal{.position = Vector2{.x = agentInit.goalPosition.x - ctx.singleton().mean.x, .y = agentInit.goalPosition.y - ctx.singleton().mean.y}}; - + ctx.get(agent)= Goal{.position = Vector3{ + .x = agentInit.goalPosition.x - ctx.singleton().mean.x, + .y = agentInit.goalPosition.y - ctx.singleton().mean.y, + .z = agentInit.goalPosition.z - ctx.singleton().mean.z}}; populateExpertTrajectory(ctx, agent, agentInit); //Applying custom rules @@ -146,19 +153,34 @@ static inline Entity createAgent(Engine &ctx, const MapObject &agentInit) { } static Entity makeRoadEdge(Engine &ctx, const MapRoad &roadInit, CountT j) { - const MapVector2 &p1 = roadInit.geometry[j]; - const MapVector2 &p2 = roadInit.geometry[j+1]; // This is guaranteed to be within bounds - - float z = 1 + (roadInit.type == EntityType::RoadEdge ? consts::lidarRoadEdgeOffset : consts::lidarRoadLineOffset); + const MapVector3 &p1 = roadInit.geometry[j]; + const MapVector3 &p2 = roadInit.geometry[j+1]; // This is guaranteed to be within bounds - Vector3 start{.x = p1.x - ctx.singleton().mean.x, .y = p1.y - ctx.singleton().mean.y, .z = z}; - Vector3 end{.x = p2.x - ctx.singleton().mean.x, .y = p2.y - ctx.singleton().mean.y, .z = z}; + Vector3 start{.x = p1.x - ctx.singleton().mean.x, .y = p1.y - ctx.singleton().mean.y, .z = p1.z - ctx.singleton().mean.z}; + Vector3 end{.x = p2.x - ctx.singleton().mean.x, .y = p2.y - ctx.singleton().mean.y, .z = p2.z - ctx.singleton().mean.z}; auto road_edge = ctx.makeRenderableEntity(); ctx.get(road_edge).e = ctx.makeEntity(); - auto pos = Vector3{.x = (start.x + end.x)/2, .y = (start.y + end.y)/2, .z = z}; - auto rot = Quat::angleAxis(atan2(end.y - start.y, end.x - start.x), madrona::math::up); + auto pos = Vector3{.x = (start.x + end.x)/2, .y = (start.y + end.y)/2, .z = (start.z + end.z)/2}; + + //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 horizontalDir{dx/horizontalDist, dy/horizontalDist, 0}; + Vector3 pitchAxis = horizontalDir.cross(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); @@ -166,13 +188,13 @@ static Entity makeRoadEdge(Engine &ctx, const MapRoad &roadInit, CountT j) { return road_edge; } -float calculateDistance(float x1, float y1, float x2, float y2) { - return sqrt(pow(x2 - x1, 2) + pow(y2 - y1, 2)); +float calculateDistance(float x1, float y1, float z1, float x2, float y2, float z2) { + return sqrt(pow(x2 - x1, 2) + pow(y2 - y1, 2) + pow(z2 - z1, 2)); } static Entity makeCube(Engine &ctx, const MapRoad &roadInit) { - MapVector2 points[] = { + MapVector3 points[] = { roadInit.geometry[0], roadInit.geometry[1], roadInit.geometry[2], @@ -183,9 +205,9 @@ static Entity makeCube(Engine &ctx, const MapRoad &roadInit) { float lengths[4]; for (int i = 0; i < 4; ++i) { - MapVector2 &p_start = points[i]; - MapVector2 &p_end = points[(i + 1) % 4]; // Wrap around to the first point - lengths[i] = calculateDistance(p_start.x, p_start.y, p_end.x, p_end.y); + MapVector3 &p_start = points[i]; + MapVector3 &p_end = points[(i + 1) % 4]; // Wrap around to the first point + lengths[i] = calculateDistance(p_start.x, p_start.y, p_start.z, p_end.x, p_end.y, p_end.z); } int maxLength_i = 0; @@ -197,8 +219,8 @@ static Entity makeCube(Engine &ctx, const MapRoad &roadInit) { minLength_i = i; } - MapVector2 &start = points[maxLength_i]; - MapVector2 &end = points[(maxLength_i + 1) % 4]; + MapVector3 &start = points[maxLength_i]; + MapVector3 &end = points[(maxLength_i + 1) % 4]; // Calculate rotation angle (assuming longer side is used to calculate angle) float angle = atan2(end.y - start.y, end.x - start.x); @@ -208,13 +230,15 @@ static Entity makeCube(Engine &ctx, const MapRoad &roadInit) { float sum_x = 0.0f; float sum_y = 0.0f; + float sum_z = 0.0f; for (const auto& point : points) { sum_x += point.x; sum_y += point.y; + sum_z += point.z; } - auto pos = Vector3{.x = sum_x/4 - ctx.singleton().mean.x, .y = sum_y/4 - ctx.singleton().mean.y, .z = 1 + consts::lidarRoadLineOffset}; + auto pos = Vector3{.x = sum_x/4 - ctx.singleton().mean.x, .y = sum_y/4 - ctx.singleton().mean.y, .z = sum_z/4 - ctx.singleton().mean.z}; auto rot = Quat::angleAxis(angle, madrona::math::up); auto scale = Diag3x3{.d0 = lengths[maxLength_i]/2, .d1 = lengths[minLength_i]/2, .d2 = 0.1}; setRoadEntitiesProps(ctx, speed_bump, pos, rot, scale, roadInit.type, ObjectID{(int32_t)SimObject::SpeedBump}, ResponseType::Static, roadInit.id, roadInit.mapType); @@ -225,11 +249,12 @@ static Entity makeCube(Engine &ctx, const MapRoad &roadInit) { static Entity makeStopSign(Engine &ctx, const MapRoad &roadInit) { float x1 = roadInit.geometry[0].x; float y1 = roadInit.geometry[0].y; + float z1 = roadInit.geometry[0].z; auto stop_sign = ctx.makeRenderableEntity(); ctx.get(stop_sign).e = ctx.makeEntity(); - auto pos = Vector3{.x = x1 - ctx.singleton().mean.x, .y = y1 - ctx.singleton().mean.y, .z = 1}; + auto pos = Vector3{.x = x1 - ctx.singleton().mean.x, .y = y1 - ctx.singleton().mean.y, .z = z1 - ctx.singleton().mean.z}; auto rot = Quat::angleAxis(0, madrona::math::up); auto scale = Diag3x3{.d0 = 0.2, .d1 = 0.2, .d2 = 1}; setRoadEntitiesProps(ctx, stop_sign, pos, rot, scale, EntityType::StopSign, ObjectID{(int32_t)SimObject::StopSign}, ResponseType::Static, roadInit.id, roadInit.mapType); @@ -359,7 +384,7 @@ void createPersistentEntities(Engine &ctx) { ctx.singleton().reset = 0; auto& means = ctx.singleton().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) { diff --git a/src/level_gen.hpp b/src/level_gen.hpp index 7336f7e6f..ca710eb65 100644 --- a/src/level_gen.hpp +++ b/src/level_gen.hpp @@ -56,7 +56,7 @@ void destroyWorld(Engine &ctx); ctx.get(road) = responseType; ctx.get(road).id = roadIdx; ctx.get(road) = mapType; - ctx.get(ctx.get(road).e) = MapObservation{.position = pos.xy(), + ctx.get(ctx.get(road).e) = MapObservation{.position = pos, .scale = scale, .heading = utils::quatToYaw(rot), .type = (float)type, diff --git a/src/sim.cpp b/src/sim.cpp index d53c167b6..19d732621 100755 --- a/src/sim.cpp +++ b/src/sim.cpp @@ -163,8 +163,8 @@ inline void collectSelfObsSystem(Engine &ctx, auto &self_obs = ctx.get(agent_iface.e); self_obs.speed = vel.linear.length(); self_obs.vehicle_size = size; - auto goalPos = goal.position - pos.xy(); - self_obs.goal.position = rot.inv().rotateVec({goalPos.x, goalPos.y, 0}).xy(); + auto goalPos = goal.position - pos; + self_obs.goal.position = rot.inv().rotateVec({goalPos.x, goalPos.y, goalPos.z}); auto hasCollided = collisionEvent.hasCollided.load_relaxed(); self_obs.collisionState = hasCollided ? 1.f : 0.f; @@ -192,8 +192,8 @@ inline void collectPartnerObsSystem(Engine &ctx, const Rotation &other_rot = ctx.get(other); const VehicleSize &other_size = ctx.get(other); - Vector2 relative_pos = (other_position - pos).xy(); - relative_pos = rot.inv().rotateVec({relative_pos.x, relative_pos.y, 0}).xy(); + Vector3 relative_pos = (other_position - pos); + relative_pos = rot.inv().rotateVec(relative_pos); float relative_speed = other_velocity.linear.length(); // Design decision: return the speed of the other agent directly Rotation relative_orientation = rot.inv() * other_rot; @@ -231,13 +231,13 @@ inline void collectMapObservationsSystem(Engine &ctx, const auto alg = ctx.data().params.roadObservationAlgorithm; if (alg == FindRoadObservationsWith::KNearestEntitiesWithRadiusFiltering) { selectKNearestRoadEntities( - ctx, rot, pos.xy(), map_obs.obs); + ctx, rot, pos, map_obs.obs); return; } assert(alg == FindRoadObservationsWith::AllEntitiesWithRadiusFiltering); - utils::ReferenceFrame referenceFrame(pos.xy(), rot); + utils::ReferenceFrame referenceFrame(pos, rot); CountT arrIndex = 0; CountT roadIdx = 0; while(roadIdx < ctx.data().numRoads && arrIndex < consts::kMaxAgentMapObservationsCount) { Entity road = ctx.data().roads[roadIdx++]; @@ -452,13 +452,13 @@ inline void rewardSystem(Engine &ctx, const auto &rewardType = ctx.data().params.rewardParams.rewardType; if(rewardType == RewardType::DistanceBased) { - float dist = (position.xy() - goal.position).length(); + float dist = (position - goal.position).length(); float reward = -dist; out_reward.v = reward; } else if(rewardType == RewardType::OnGoalAchieved) { - float dist = (position.xy() - goal.position).length(); + float dist = (position - goal.position).length(); float reward = (dist < ctx.data().params.rewardParams.distanceToGoalThreshold) ? 1.f : 0.f; out_reward.v = reward; } @@ -502,7 +502,7 @@ inline void doneSystem(Engine &ctx, // An agent can be done early if it reaches the goal if (done.v != 1 || info.reachedGoal != 1) { - float dist = (position.xy() - goal.position).length(); + float dist = (position - goal.position).length(); if (dist < ctx.data().params.rewardParams.distanceToGoalThreshold) { done.v = 1; diff --git a/src/types.hpp b/src/types.hpp index 13dc4dd47..fbf9c8356 100755 --- a/src/types.hpp +++ b/src/types.hpp @@ -76,7 +76,7 @@ struct AgentID { struct Goal { - madrona::math::Vector2 position; + madrona::math::Vector3 position; }; // WorldReset is a per-world singleton component that causes the current @@ -93,7 +93,7 @@ struct AgentID { }; struct WorldMeans { - madrona::math::Vector3 mean; // TODO: Z is 0 for now, but can be used for 3D in future + madrona::math::Vector3 mean; }; const size_t WorldMeansExportSize = 3; @@ -185,19 +185,19 @@ struct AgentID { return SelfObservation{ .speed = 0, .vehicle_size = {0, 0}, - .goal = {.position = {0, 0}}, + .goal = {.position = {0, 0, 0}}, .collisionState = 0, .id = -1}; } }; - const size_t SelfObservationExportSize = 7; + const size_t SelfObservationExportSize = 8; static_assert(sizeof(SelfObservation) == sizeof(float) * SelfObservationExportSize); struct MapObservation { - madrona::math::Vector2 position; + madrona::math::Vector3 position; Scale scale; float heading; float type; @@ -207,7 +207,7 @@ struct AgentID { static inline MapObservation zero() { return MapObservation{ - .position = {0, 0}, + .position = {0, 0, 0}, .scale = madrona::math::Diag3x3{0, 0, 0}, .heading = 0, .type = static_cast(EntityType::None), @@ -217,14 +217,14 @@ struct AgentID { } }; - const size_t MapObservationExportSize = 9; + const size_t MapObservationExportSize = 10; static_assert(sizeof(MapObservation) == sizeof(float) * MapObservationExportSize); struct PartnerObservation { float speed; - madrona::math::Vector2 position; + madrona::math::Vector3 position; float heading; VehicleSize vehicle_size; float type; @@ -233,7 +233,7 @@ struct AgentID { static inline PartnerObservation zero() { return PartnerObservation{ .speed = 0, - .position = {0, 0}, + .position = {0, 0, 0}, .heading = 0, .vehicle_size = {0, 0}, .type = static_cast(EntityType::None), @@ -255,7 +255,7 @@ struct AgentID { PartnerObservation obs[consts::kMaxAgentCount - 1]; }; - const size_t PartnerObservationExportSize = 8; + const size_t PartnerObservationExportSize = 9; static_assert(sizeof(PartnerObservations) == sizeof(float) * (consts::kMaxAgentCount - 1) * PartnerObservationExportSize); @@ -275,7 +275,7 @@ struct AgentID { { float depth; float encodedType; - madrona::math::Vector2 position; + madrona::math::Vector3 position; }; // Linear depth values and entity type in a circle around the agent @@ -286,7 +286,7 @@ struct AgentID { LidarSample samplesRoadLines[consts::numLidarSamples]; }; - const size_t LidarExportSize = 3 * consts::numLidarSamples * 4; + const size_t LidarExportSize = 3 * consts::numLidarSamples * 5; static_assert(sizeof(Lidar) == sizeof(float) * LidarExportSize); // Number of steps remaining in the episode. Allows non-recurrent policies @@ -311,14 +311,14 @@ struct AgentID { struct Trajectory { - madrona::math::Vector2 positions[consts::kTrajectoryLength]; + madrona::math::Vector3 positions[consts::kTrajectoryLength]; madrona::math::Vector2 velocities[consts::kTrajectoryLength]; float headings[consts::kTrajectoryLength]; float valids[consts::kTrajectoryLength]; Action inverseActions[consts::kTrajectoryLength]; }; - const size_t TrajectoryExportSize = 2 * 2 * consts::kTrajectoryLength + 2 * consts::kTrajectoryLength + ActionExportSize * consts::kTrajectoryLength; + const size_t TrajectoryExportSize = 3 * consts::kTrajectoryLength + 2 * 2 * consts::kTrajectoryLength + ActionExportSize * consts::kTrajectoryLength; static_assert(sizeof(Trajectory) == sizeof(float) * TrajectoryExportSize); @@ -353,7 +353,7 @@ struct AgentID { float id; }; - const size_t AbsoluteSelfObservationExportSize = 13; // 3 + 4 + 1 + 2 + 2 + const size_t AbsoluteSelfObservationExportSize = 14; // 3 + 5 + 3 + 2 + 1 static_assert(sizeof(AbsoluteSelfObservation) == sizeof(float) * AbsoluteSelfObservationExportSize); diff --git a/src/utils.hpp b/src/utils.hpp index 10751e082..406c5ca9e 100644 --- a/src/utils.hpp +++ b/src/utils.hpp @@ -26,7 +26,7 @@ inline float quatToYaw(madrona::base::Rotation q) { class ReferenceFrame { public: - ReferenceFrame(const madrona::math::Vector2 &position, + ReferenceFrame(const madrona::math::Vector3 &position, const madrona::base::Rotation &rotation) : referenceRotation(rotation), referencePosition(position) {} @@ -47,20 +47,19 @@ class ReferenceFrame { } private: - madrona::math::Vector2 + madrona::math::Vector3 relative(const madrona::math::Vector3 &absolutePos) const { - auto relativePosition = absolutePos.xy() - referencePosition; + auto relativePosition = absolutePos - referencePosition; return referenceRotation.inv() - .rotateVec({relativePosition.x, relativePosition.y, 0}) - .xy(); + .rotateVec({relativePosition.x, relativePosition.y, relativePosition.z}); } float relative(const madrona::base::Rotation &absoluteRot) const { return gpudrive::utils::quatToYaw(referenceRotation.inv() * absoluteRot); } - madrona::math::Vector2 referencePosition; + madrona::math::Vector3 referencePosition; madrona::base::Rotation referenceRotation; }; }} diff --git a/tests/EgocentricRoadObservationTests.cpp b/tests/EgocentricRoadObservationTests.cpp index 32701b5d0..7c70307f6 100644 --- a/tests/EgocentricRoadObservationTests.cpp +++ b/tests/EgocentricRoadObservationTests.cpp @@ -7,16 +7,17 @@ using namespace madrona::math; using gpudrive::utils::ReferenceFrame; TEST(EgocentricRoadObservationTests, Relative) { - ReferenceFrame rf{Vector2{.x = 3, .y = 0}, + ReferenceFrame rf{Vector3{.x = 3, .y = 0, .z=0}, Quat::angleAxis(toRadians(90), madrona::math::up)}; auto mapObs = - rf.observationOf(Vector3{.x = 3, .y = 3, .z = 0}, + rf.observationOf(Vector3{.x = 3, .y = 3, .z = 3}, Quat::angleAxis(toRadians(270), madrona::math::up), Scale{{.d0 = 10, .d1 = 0.1, .d2 = 0.1}}, static_cast(0), 0); EXPECT_LT(mapObs.position.x - 3, 0.000001); EXPECT_LT(mapObs.position.y - 0, 0.000001); + EXPECT_LT(mapObs.position.z - 3, 0.000001); EXPECT_EQ(mapObs.heading, toRadians(180)); }