Skip to content

Commit

Permalink
Remove padding agents using interface entities (#103)
Browse files Browse the repository at this point in the history 
* Init remove padding agents

* Testing

* Refactor out the controlled state

* Fix the tests

* Use by reference

* Temp fix for map export

* cleanup

* Fix merge issues

* Fix merge issues

* make consts 6000 to pass tests

* Use init counts for BVH

* Remove unused function

* Pass all tests

* Set controlledstate

* Zero out things

* Zero padded agents out

* Add a test.py for easy debugging

* Better check for map drawing.

* Set agents to 128

* Cycle through files in test.py

* Rename InterfaceEntity to AgentInterfaceEntity

* Only print dones

* Remove debug checks

* Separete out the observation systems

* Sort interfaces

* remove debug script

* Remove debug statements

* Minor improvements

* Info tensor interface matches main

* Correctly export response type and trajectory
  • Loading branch information
@aaravpandya
aaravpandya authored Sep 21, 2024
1 parent 8a59d44 commit dc160be
Show file tree
Hide file tree
Showing 12 changed files with 307 additions and 266 deletions.
2 changes: 1 addition & 1 deletion pygpudrive/env/viz.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -277,7 +277,7 @@ def draw_map(self, surf, map_info, world_render_idx=0):
"""Draw static map elements."""
for idx, map_obj in enumerate(map_info):

if map_obj[-1] == float(gpudrive.EntityType.Padding):
if map_obj[-1] == float(gpudrive.EntityType.Padding) or map_obj[-1] == float(gpudrive.EntityType._None):
continue

elif map_obj[-1] <= float(gpudrive.EntityType.RoadLane):
Expand Down
2 changes: 1 addition & 1 deletion src/consts.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
#include <limits>
#include <madrona/math.hpp>
#include <madrona/types.hpp>

namespace gpudrive {

namespace consts {
Expand Down
2 changes: 1 addition & 1 deletion src/dynamics.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ using namespace madrona::math;
namespace gpudrive
{

inline void forwardKinematics(Action &action, VehicleSize &size, Rotation &rotation, Position &position, Velocity &velocity)
inline void forwardKinematics(const Action &action, VehicleSize &size, Rotation &rotation, Position &position, Velocity &velocity)
{
const float maxSpeed{std::numeric_limits<float>::max()};
const float dt{0.1};
Expand Down
13 changes: 7 additions & 6 deletions src/headless.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -97,17 +97,17 @@ int main(int argc, char *argv[])
// printf("Reward\n");
// rewardPrinter.print();

// printf("Done\n");
// donePrinter.print();
printf("Done\n");
donePrinter.print();

// printf("Controlled State\n");
// controlledStatePrinter.print();

printf("Agent Map Obs\n");
agent_map_obs_printer.print();
// printf("Agent Map Obs\n");
// agent_map_obs_printer.print();

printf("Info\n");
info_printer.print();
// printf("Info\n");
// info_printer.print();
};

auto worldToShape =
Expand All @@ -133,6 +133,7 @@ int main(int argc, char *argv[])
}
const auto end = std::chrono::steady_clock::now();
const std::chrono::duration<double> elapsed = end - start;
printObs();

float fps = (double)num_steps * (double)num_worlds / elapsed.count();
printf("FPS %f\n", fps);
Expand Down
3 changes: 2 additions & 1 deletion src/knn.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,8 +60,9 @@ bool isObservationsValid(gpudrive::Engine &ctx,
sortedObservations.reserve(roadCount);

for (madrona::CountT roadIdx = 0; roadIdx < roadCount; ++roadIdx) {
auto &road_iface = ctx.get<gpudrive::RoadInterfaceEntity>(ctx.data().roads[roadIdx]).e;
const auto &currentObservation =
ctx.get<gpudrive::MapObservation>(ctx.data().roads[roadIdx]);
ctx.get<gpudrive::MapObservation>(road_iface);
sortedObservations.emplace_back(relativeObservation(
currentObservation, referenceRotation, referencePosition));
}
Expand Down
184 changes: 79 additions & 105 deletions src/level_gen.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,12 +22,13 @@ static void registerRigidBodyEntity(
}

static inline void resetAgent(Engine &ctx, Entity agent) {
auto xCoord = ctx.get<Trajectory>(agent).positions[0].x;
auto yCoord = ctx.get<Trajectory>(agent).positions[0].y;
auto xVelocity = ctx.get<Trajectory>(agent).velocities[0].x;
auto yVelocity = ctx.get<Trajectory>(agent).velocities[0].y;
auto speed = ctx.get<Trajectory>(agent).velocities[0].length();
auto heading = ctx.get<Trajectory>(agent).headings[0];
auto agent_iface = ctx.get<AgentInterfaceEntity>(agent).e;
auto xCoord = ctx.get<Trajectory>(agent_iface).positions[0].x;
auto yCoord = ctx.get<Trajectory>(agent_iface).positions[0].y;
auto xVelocity = ctx.get<Trajectory>(agent_iface).velocities[0].x;
auto yVelocity = ctx.get<Trajectory>(agent_iface).velocities[0].y;
auto speed = ctx.get<Trajectory>(agent_iface).velocities[0].length();
auto heading = ctx.get<Trajectory>(agent_iface).headings[0];

ctx.get<Position>(agent) = Vector3{.x = xCoord, .y = yCoord, .z = 1};
ctx.get<Rotation>(agent) = Quat::angleAxis(heading, madrona::math::up);
Expand All @@ -37,30 +38,31 @@ static inline void resetAgent(Engine &ctx, Entity agent) {
{
case DynamicsModel::Classic:
{
ctx.get<Action>(agent) = Action{.classic = {0, 0, 0}};
ctx.get<Action>(agent_iface) = Action{.classic = {0, 0, 0}};
break;
}
case DynamicsModel::InvertibleBicycle:
{
ctx.get<Action>(agent) = Action{.classic = {0, 0, 0}};
ctx.get<Action>(agent_iface) = Action{.classic = {0, 0, 0}};
break;
}
case DynamicsModel::DeltaLocal:
{
ctx.get<Action>(agent) = Action{.delta{.dx = 0, .dy = 0, .dyaw = 0}};
ctx.get<Action>(agent_iface) = Action{.delta{.dx = 0, .dy = 0, .dyaw = 0}};
break;
}
case DynamicsModel::State:
{
ctx.get<Action>(agent) = Action{.state = {.position = Vector3{0, 0, 1}, .yaw = 0, .velocity = {.linear = Vector3::zero(), .angular = Vector3::zero()}}};
ctx.get<Action>(agent_iface) = Action{.state = {.position = Vector3{0, 0, 1}, .yaw = 0, .velocity = {.linear = Vector3::zero(), .angular = Vector3::zero()}}};
break;
}
}
ctx.get<StepsRemaining>(agent).t = consts::episodeLen;
ctx.get<Done>(agent).v = 0;
ctx.get<Reward>(agent).v = 0;
ctx.get<Info>(agent) = Info{};
ctx.get<Info>(agent).type = (int32_t)ctx.get<EntityType>(agent);
ctx.get<StepsRemaining>(agent_iface).t = consts::episodeLen;
ctx.get<Done>(agent_iface).v = 0;
ctx.get<Reward>(agent_iface).v = 0;
ctx.get<Info>(agent_iface) = Info{};
ctx.get<Info>(agent_iface).type = (int32_t)ctx.get<EntityType>(agent);
ctx.get<ResponseType>(agent_iface) = ctx.get<ResponseType>(agent);

if(ctx.get<ResponseType>(agent) == ResponseType::Static)
{
Expand All @@ -74,7 +76,8 @@ static inline void resetAgent(Engine &ctx, Entity agent) {
}

static inline void populateExpertTrajectory(Engine &ctx, const Entity &agent, const MapObject &agentInit) {
auto &trajectory = ctx.get<Trajectory>(agent);
const auto &agent_iface = ctx.get<AgentInterfaceEntity>(agent).e;
auto &trajectory = ctx.get<Trajectory>(agent_iface);
for(CountT i = 0; i < agentInit.numPositions; i++)
{
trajectory.positions[i] = Vector2{.x = agentInit.position[i].x - ctx.data().mean.x, .y = agentInit.position[i].y - ctx.data().mean.y};
Expand Down Expand Up @@ -141,21 +144,23 @@ static inline Entity createAgent(Engine &ctx, const MapObject &agentInit) {
assert(agentInit.type >= EntityType::Vehicle || agentInit.type == EntityType::None);
ctx.get<EntityType>(agent) = agentInit.type;

auto agent_iface = ctx.get<AgentInterfaceEntity>(agent).e = ctx.makeEntity<AgentInterface>();

ctx.get<Goal>(agent)= Goal{.position = Vector2{.x = agentInit.goalPosition.x - ctx.data().mean.x, .y = agentInit.goalPosition.y - ctx.data().mean.y}};
populateExpertTrajectory(ctx, agent, agentInit);
if(!ctx.data().params.isStaticAgentControlled && (ctx.get<Goal>(agent).position - ctx.get<Trajectory>(agent).positions[0]).length() < consts::staticThreshold)
if(!ctx.data().params.isStaticAgentControlled && (ctx.get<Goal>(agent).position - ctx.get<Trajectory>(agent_iface).positions[0]).length() < consts::staticThreshold)
{
ctx.get<ResponseType>(agent) = ResponseType::Static;
}

if(ctx.data().numControlledVehicles < ctx.data().params.maxNumControlledVehicles && agentInit.type == EntityType::Vehicle && agentInit.valid[0] && ctx.get<ResponseType>(agent) == ResponseType::Dynamic)
{
ctx.get<ControlledState>(agent) = ControlledState{.controlled = 1};
ctx.get<ControlledState>(agent_iface) = ControlledState{.controlled = 1};
ctx.data().numControlledVehicles++;
}
else
{
ctx.get<ControlledState>(agent) = ControlledState{.controlled = 0};
ctx.get<ControlledState>(agent_iface) = ControlledState{.controlled = 0};
}

// This is not stricly necessary since , but is kept here for consistency
Expand Down Expand Up @@ -189,7 +194,8 @@ static Entity makeRoadEdge(Engine &ctx, const MapVector2 &p1,
ctx.get<ObjectID>(road_edge) = ObjectID{(int32_t)SimObject::Cube};
registerRigidBodyEntity(ctx, road_edge, SimObject::Cube);
ctx.get<ResponseType>(road_edge) = ResponseType::Static;
ctx.get<MapObservation>(road_edge) = MapObservation{.position = ctx.get<Position>(road_edge).xy(),
auto road_iface = ctx.get<RoadInterfaceEntity>(road_edge).e = ctx.makeEntity<RoadInterface>();
ctx.get<MapObservation>(road_iface) = MapObservation{.position = ctx.get<Position>(road_edge).xy(),
.scale = ctx.get<Scale>(road_edge),
.heading = utils::quatToYaw(ctx.get<Rotation>(road_edge)),
.type = (float)type};
Expand Down Expand Up @@ -262,7 +268,8 @@ static Entity makeCube(Engine &ctx, const MapVector2 &p1, const MapVector2 &p2,
ctx.get<ObjectID>(speed_bump) = ObjectID{(int32_t)SimObject::SpeedBump};
registerRigidBodyEntity(ctx, speed_bump, SimObject::SpeedBump);
ctx.get<ResponseType>(speed_bump) = ResponseType::Static;
ctx.get<MapObservation>(speed_bump) = MapObservation{.position = ctx.get<Position>(speed_bump).xy(),
auto road_iface = ctx.get<RoadInterfaceEntity>(speed_bump).e = ctx.makeEntity<RoadInterface>();
ctx.get<MapObservation>(road_iface) = MapObservation{.position = ctx.get<Position>(speed_bump).xy(),
.scale = ctx.get<Scale>(speed_bump),
.heading = utils::quatToYaw(ctx.get<Rotation>(speed_bump)),
.type = (float)type};
Expand All @@ -281,7 +288,8 @@ static Entity makeStopSign(Engine &ctx, const MapVector2 &p1) {
ctx.get<ObjectID>(stop_sign) = ObjectID{(int32_t)SimObject::StopSign};
registerRigidBodyEntity(ctx, stop_sign, SimObject::StopSign);
ctx.get<ResponseType>(stop_sign) = ResponseType::Static;
ctx.get<MapObservation>(stop_sign) = MapObservation{.position = ctx.get<Position>(stop_sign).xy(),
auto road_iface = ctx.get<RoadInterfaceEntity>(stop_sign).e = ctx.makeEntity<RoadInterface>();
ctx.get<MapObservation>(road_iface) = MapObservation{.position = ctx.get<Position>(stop_sign).xy(),
.scale = ctx.get<Scale>(stop_sign),
.heading = utils::quatToYaw(ctx.get<Rotation>(stop_sign)),
.type = (float)EntityType::StopSign};
Expand All @@ -292,27 +300,36 @@ static inline void createRoadEntities(Engine &ctx, const MapRoad &roadInit, Coun
if (roadInit.type == EntityType::RoadEdge || roadInit.type == EntityType::RoadLine || roadInit.type == EntityType::RoadLane)
{
size_t numPoints = roadInit.numPoints;
for(size_t j = 1; j <= numPoints - 1; j++)
for (size_t j = 1; j <= numPoints - 1; j++)
{
if(idx >= consts::kMaxRoadEntityCount)
return;
ctx.data().roads[idx++] = makeRoadEdge(ctx, roadInit.geometry[j-1], roadInit.geometry[j], roadInit.type);
if (idx >= consts::kMaxRoadEntityCount)
return;
auto road = ctx.data().roads[idx] = makeRoadEdge(ctx, roadInit.geometry[j - 1], roadInit.geometry[j], roadInit.type);
ctx.data().road_ifaces[idx++] = ctx.get<RoadInterfaceEntity>(road).e;
}
} else if (roadInit.type == EntityType::SpeedBump || roadInit.type == EntityType::CrossWalk) {
assert(roadInit.numPoints >= 4);
// TODO: Speed Bump are not guranteed to have 4 points. Need to handle this case.
if(idx >= consts::kMaxRoadEntityCount)
return;
ctx.data().roads[idx++] = makeCube(ctx, roadInit.geometry[0], roadInit.geometry[1], roadInit.geometry[2], roadInit.geometry[3], roadInit.type);
} else if (roadInit.type == EntityType::StopSign ) {
assert(roadInit.numPoints >= 1);
// TODO: Stop Sign are not guranteed to have 1 point. Need to handle this case.
if(idx >= consts::kMaxRoadEntityCount)
return;
ctx.data().roads[idx++] = makeStopSign(ctx, roadInit.geometry[0]);
} else {
// TODO: Need to handle Cross Walk.
// assert(false);
}
else if (roadInit.type == EntityType::SpeedBump || roadInit.type == EntityType::CrossWalk)
{
assert(roadInit.numPoints >= 4);
// TODO: Speed Bump are not guranteed to have 4 points. Need to handle this case.
if (idx >= consts::kMaxRoadEntityCount)
return;
auto road = ctx.data().roads[idx] = makeCube(ctx, roadInit.geometry[0], roadInit.geometry[1], roadInit.geometry[2], roadInit.geometry[3], roadInit.type);
ctx.data().road_ifaces[idx++] = ctx.get<RoadInterfaceEntity>(road).e;
}
else if (roadInit.type == EntityType::StopSign)
{
assert(roadInit.numPoints >= 1);
// TODO: Stop Sign are not guranteed to have 1 point. Need to handle this case.
if (idx >= consts::kMaxRoadEntityCount)
return;
auto road = ctx.data().roads[idx] = makeStopSign(ctx, roadInit.geometry[0]);
ctx.data().road_ifaces[idx++] = ctx.get<RoadInterfaceEntity>(road).e;
}
else
{
// TODO: Need to handle Cross Walk.
// assert(false);
return;
}
}
Expand All @@ -330,58 +347,33 @@ static void createFloorPlane(Engine &ctx)
registerRigidBodyEntity(ctx, ctx.data().floorPlane, SimObject::Plane);
}

static inline Entity createAgentPadding(Engine &ctx) {
auto agent = ctx.makeRenderableEntity<Agent>();

ctx.get<Position>(agent) = consts::kPaddingPosition;
ctx.get<Rotation>(agent) = Quat::angleAxis(0, madrona::math::up);
ctx.get<Scale>(agent) = Diag3x3{.d0 = 0, .d1 = 0, .d2 = 0};
ctx.get<Velocity>(agent) = {Vector3::zero(), Vector3::zero()};
ctx.get<ObjectID>(agent) = ObjectID{(int32_t)SimObject::Agent};
ctx.get<ResponseType>(agent) = ResponseType::Static;
ctx.get<EntityType>(agent) = EntityType::Padding;
ctx.get<CollisionDetectionEvent>(agent).hasCollided.store_release(0);
ctx.get<Done>(agent).v = 0;
ctx.get<StepsRemaining>(agent).t = consts::episodeLen;
ctx.get<ControlledState>(agent) = ControlledState{.controlled = 0};

if (ctx.data().enableRender) {
render::RenderingSystem::attachEntityToView(ctx,
agent,
90.f, 0.001f,
1.5f * math::up);
}

return agent;
}

static inline Entity createPhysicsEntityPadding(Engine &ctx) {
auto physicsEntity = ctx.makeRenderableEntity<PhysicsEntity>();

ctx.get<Position>(physicsEntity) = consts::kPaddingPosition;
ctx.get<Rotation>(physicsEntity) = Quat::angleAxis(0, madrona::math::up);
ctx.get<Scale>(physicsEntity) = Diag3x3{.d0 = 0, .d1 = 0, .d2 = 0};
ctx.get<Velocity>(physicsEntity) = {Vector3::zero(), Vector3::zero()};
ctx.get<ObjectID>(physicsEntity) = ObjectID{(int32_t)SimObject::Cube};
ctx.get<ResponseType>(physicsEntity) = ResponseType::Static;
ctx.get<MapObservation>(physicsEntity) = MapObservation{.position = ctx.get<Position>(physicsEntity).xy(),
.scale = ctx.get<Scale>(physicsEntity),
.heading = utils::quatToYaw(ctx.get<Rotation>(physicsEntity)),
.type = float(EntityType::Padding)};
ctx.get<EntityType>(physicsEntity) = EntityType::Padding;

return physicsEntity;
}

void createPaddingEntities(Engine &ctx) {
for (CountT agentIdx = ctx.data().numAgents;
agentIdx < consts::kMaxAgentCount; ++agentIdx) {
ctx.data().agents[agentIdx] = createAgentPadding(ctx);
Entity &agent_iface = ctx.data().agent_ifaces[agentIdx] = ctx.makeEntity<AgentInterface>();
ctx.get<ControlledState>(agent_iface) = ControlledState{.controlled = 0};
ctx.get<Done>(agent_iface).v = 1;
ctx.get<Reward>(agent_iface).v = 0;
ctx.get<Info>(agent_iface) = Info::zero();
ctx.get<ResponseType>(agent_iface) = ResponseType::Static;
auto &agent_map_obs = ctx.get<AgentMapObservations>(agent_iface);
for (CountT i = 0; i < consts::kMaxAgentMapObservationsCount; i++) {
agent_map_obs.obs[i] = MapObservation::zero();
}
auto &self_obs = ctx.get<SelfObservation>(agent_iface);
self_obs = SelfObservation::zero();

auto &partner_obs = ctx.get<PartnerObservations>(agent_iface);
for (CountT i = 0; i < consts::kMaxAgentCount-1; i++) {
partner_obs.obs[i] = PartnerObservation::zero();
}

}

for (CountT roadIdx = ctx.data().numRoads;
roadIdx < consts::kMaxRoadEntityCount; ++roadIdx) {
ctx.data().roads[roadIdx] = createPhysicsEntityPadding(ctx);
Entity &e = ctx.data().road_ifaces[roadIdx] = ctx.makeEntity<RoadInterface>();
ctx.get<MapObservation>(e) = MapObservation::zero();
}
}

Expand Down Expand Up @@ -429,6 +421,7 @@ void createPersistentEntities(Engine &ctx, Map *map) {
}
auto agent = createAgent(
ctx, agentInit);
ctx.data().agent_ifaces[agentIdx] = ctx.get<AgentInterfaceEntity>(agent).e;
ctx.data().agents[agentIdx++] = agent;
}

Expand All @@ -451,24 +444,6 @@ void createPersistentEntities(Engine &ctx, Map *map) {
createPaddingEntities(ctx);
}

static void resetPaddingEntities(Engine &ctx) {
for (CountT agentIdx = ctx.data().numAgents;
agentIdx < consts::kMaxAgentCount; ++agentIdx) {
Entity agent = ctx.data().agents[agentIdx];
ctx.get<Done>(agent).v = 0;
ctx.get<StepsRemaining>(agent).t = consts::episodeLen;
ctx.get<Info>(agent) = Info{};
ctx.get<Info>(agent).type = (int32_t)ctx.get<EntityType>(agent);
registerRigidBodyEntity(ctx, agent, SimObject::Agent);
}

for (CountT roadIdx = ctx.data().numRoads;
roadIdx < consts::kMaxRoadEntityCount; ++roadIdx) {
Entity road = ctx.data().roads[roadIdx];
registerRigidBodyEntity(ctx, road, SimObject::Cube);
}
}

static void resetPersistentEntities(Engine &ctx)
{

Expand Down Expand Up @@ -516,7 +491,6 @@ static void resetPersistentEntities(Engine &ctx)
void generateWorld(Engine &ctx)
{
resetPersistentEntities(ctx);
resetPaddingEntities(ctx);
}

}
2 changes: 1 addition & 1 deletion src/mgr.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -676,7 +676,7 @@ Tensor Manager::agentMapObservationsTensor() const
TensorElementType::Float32,
{
impl_->numWorlds,
consts::kMaxAgentCount,
consts::kMaxAgentCount,
consts::kMaxAgentMapObservationsCount,
AgentMapObservationExportSize,
});
Expand Down
Loading

0 comments on commit dc160be

Please sign in to comment.