From 0c65bfb4344a383a49595d088f0a23c6cc1d21be Mon Sep 17 00:00:00 2001 From: Aarav Pandya Date: Sun, 6 Apr 2025 23:10:28 -0400 Subject: [PATCH] init --- src/bindings.cpp | 5 +++-- src/consts.hpp | 4 ++-- src/init.hpp | 3 ++- src/json_serialization.hpp | 9 +++++---- src/level_gen.cpp | 4 ++-- src/mgr.cpp | 6 ++++++ src/mgr.hpp | 1 + src/sim.cpp | 15 ++++++++++----- src/sim.hpp | 1 + src/types.hpp | 19 ++++++++++++------- src/viewer.cpp | 2 +- 11 files changed, 45 insertions(+), 24 deletions(-) diff --git a/src/bindings.cpp b/src/bindings.cpp index 6b7414ddb..dd4a75c23 100755 --- a/src/bindings.cpp +++ b/src/bindings.cpp @@ -23,7 +23,7 @@ namespace madrona_gpudrive m.attr("kMaxAgentCount") = consts::kMaxAgentCount; m.attr("kMaxRoadEntityCount") = consts::kMaxRoadEntityCount; m.attr("kMaxAgentMapObservationsCount") = consts::kMaxAgentMapObservationsCount; - m.attr("episodeLen") = consts::episodeLen; + m.attr("maxEpisodeLen") = consts::maxEpisodeLength; m.attr("numLidarSamples") = consts::numLidarSamples; m.attr("vehicleScale") = consts::vehicleLengthScale; @@ -146,7 +146,8 @@ namespace madrona_gpudrive }) .def("deleted_agents_tensor", &Manager::deletedAgentsTensor) .def("map_name_tensor", &Manager::mapNameTensor) - .def("scenario_id_tensor", &Manager::scenarioIdTensor); + .def("scenario_id_tensor", &Manager::scenarioIdTensor) + .def("episode_length_tensor", &Manager::episodeLengthTensor); } } diff --git a/src/consts.hpp b/src/consts.hpp index 41d7a5282..7f2c13a07 100644 --- a/src/consts.hpp +++ b/src/consts.hpp @@ -31,7 +31,7 @@ inline constexpr float rewardPerDist = 0.05f; inline constexpr float slackReward = -0.005f; // Steps per episode -inline constexpr int32_t episodeLen = 91; +inline constexpr int32_t maxEpisodeLength = 200; // Number of lidar samples, arranged in circle around agent inline constexpr madrona::CountT numLidarSamples = 50; @@ -57,7 +57,7 @@ inline constexpr madrona::CountT numPhysicsSubsteps = 0.f; inline constexpr float zDimensionScale = 1; inline constexpr float xDimensionScaleRoadSegment = 1; -inline constexpr madrona::CountT kTrajectoryLength = 91; // Nocturne has 90 timesteps per episode. making it 91 as a buffer. +// inline constexpr madrona::CountT kTrajectoryLength = 91; // Nocturne has 90 timesteps per episode. making it 91 as a buffer. inline constexpr madrona::CountT kMaxRoadGeometryLength = 1810; diff --git a/src/init.hpp b/src/init.hpp index 253e1d06e..6c46539f8 100755 --- a/src/init.hpp +++ b/src/init.hpp @@ -8,7 +8,7 @@ namespace madrona_gpudrive // Constants computed from train files. constexpr size_t MAX_OBJECTS = 515; constexpr size_t MAX_ROADS = 956; - constexpr size_t MAX_POSITIONS = 91; + constexpr size_t MAX_POSITIONS = consts::maxEpisodeLength; constexpr size_t MAX_GEOMETRY = 1746; // Cannot use Madrona::math::Vector2 because it is not a POD type. @@ -124,6 +124,7 @@ namespace madrona_gpudrive bool disableClassicalObs = false; DynamicsModel dynamicsModel = DynamicsModel::Classic; bool readFromTracksToPredict = false; // Default: false - for womd_tracks_to_predict initialization mode + uint32_t episode_length = 91; }; struct WorldInit diff --git a/src/json_serialization.hpp b/src/json_serialization.hpp index 95d651d9a..04c3e5517 100644 --- a/src/json_serialization.hpp +++ b/src/json_serialization.hpp @@ -18,9 +18,10 @@ namespace madrona_gpudrive { obj.mean = {0,0}; uint32_t i = 0; + int numPositions = j.at("position").size(); for (const auto &pos : j.at("position")) { - if (i < MAX_POSITIONS) + if (i < MAX_POSITIONS && i < numPositions) { from_json(pos, obj.position[i]); obj.mean.x += (obj.position[i].x - obj.mean.x)/(i+1); @@ -41,7 +42,7 @@ namespace madrona_gpudrive i = 0; for (const auto &h : j.at("heading")) { - if (i < MAX_POSITIONS) + if (i < MAX_POSITIONS && i < numPositions) { h.get_to(obj.heading[i]); ++i; @@ -56,7 +57,7 @@ namespace madrona_gpudrive i = 0; for (const auto &v : j.at("velocity")) { - if (i < MAX_POSITIONS) + if (i < MAX_POSITIONS && i < numPositions) { from_json(v, obj.velocity[i]); ++i; @@ -71,7 +72,7 @@ namespace madrona_gpudrive i = 0; for (const auto &v : j.at("valid")) { - if (i < MAX_POSITIONS) + if (i < MAX_POSITIONS && i < numPositions) { v.get_to(obj.valid[i]); ++i; diff --git a/src/level_gen.cpp b/src/level_gen.cpp index e550efff1..5b7098160 100755 --- a/src/level_gen.cpp +++ b/src/level_gen.cpp @@ -20,7 +20,7 @@ static void registerRigidBodyEntity( ctx.get(e) = PhysicsSystem::registerEntity(ctx, e, obj_id); } -static inline void resetAgentInterface(Engine &ctx, Entity agent_iface, EntityType type, ResponseType resp_type, int32_t steps_remaining= consts::episodeLen, int32_t done = 0) { +static inline void resetAgentInterface(Engine &ctx, Entity agent_iface, EntityType type, ResponseType resp_type, int32_t steps_remaining, int32_t done = 0) { ctx.get(agent_iface).t = steps_remaining; ctx.get(agent_iface).v = done; ctx.get(agent_iface).v = 0; @@ -46,7 +46,7 @@ static inline void resetAgent(Engine &ctx, Entity agent) { } ctx.get(agent_iface) = getZeroAction(ctx.data().params.dynamicsModel); - resetAgentInterface(ctx, agent_iface, ctx.get(agent), ctx.get(agent)); + resetAgentInterface(ctx, agent_iface, ctx.get(agent), ctx.get(agent), ctx.singleton().episode_length); #ifndef GPUDRIVE_DISABLE_NARROW_PHASE ctx.get(agent).hasCollided.store_release(0); diff --git a/src/mgr.cpp b/src/mgr.cpp index e314735a4..13cc52bc8 100755 --- a/src/mgr.cpp +++ b/src/mgr.cpp @@ -812,6 +812,12 @@ Tensor Manager::agentMapObservationsTensor() const } +Tensor Manager::episodeLengthTensor() const +{ + return impl_->exportTensor(ExportID::EpisodeLength, TensorElementType::Int32, + {impl_->numWorlds, 1}); +} + Tensor Manager::lidarTensor() const { return impl_->exportTensor(ExportID::Lidar, TensorElementType::Float32, diff --git a/src/mgr.hpp b/src/mgr.hpp index d9613b61b..700e85232 100755 --- a/src/mgr.hpp +++ b/src/mgr.hpp @@ -73,6 +73,7 @@ class Manager { MGR_EXPORT madrona::py::Tensor deletedAgentsTensor() const; MGR_EXPORT madrona::py::Tensor mapNameTensor() const; MGR_EXPORT madrona::py::Tensor scenarioIdTensor() const; + MGR_EXPORT madrona::py::Tensor episodeLengthTensor() const; madrona::py::Tensor rgbTensor() const; madrona::py::Tensor depthTensor() const; // These functions are used by the viewer to control the simulation diff --git a/src/sim.cpp b/src/sim.cpp index bcc0fe51e..7b0b8af04 100755 --- a/src/sim.cpp +++ b/src/sim.cpp @@ -20,8 +20,8 @@ namespace RenderingSystem = madrona::render::RenderingSystem; namespace madrona_gpudrive { -CountT getCurrentStep(const StepsRemaining &stepsRemaining) { - return consts::episodeLen - stepsRemaining.t; +CountT getCurrentStep(Engine &ctx, const StepsRemaining &stepsRemaining) { + return ctx.singleton().episode_length - stepsRemaining.t; } // Register all the ECS components and archetypes that will be @@ -68,6 +68,7 @@ void Sim::registerTypes(ECSRegistry ®istry, const Config &cfg) registry.registerSingleton(); registry.registerSingleton(); registry.registerSingleton(); + registry.registerSingleton(); registry.registerArchetype(); registry.registerArchetype(); @@ -83,6 +84,7 @@ void Sim::registerTypes(ECSRegistry ®istry, const Config &cfg) registry.exportSingleton((uint32_t)ExportID::DeletedAgents); registry.exportSingleton((uint32_t)ExportID::MapName); registry.exportSingleton((uint32_t)ExportID::ScenarioId); + registry.exportSingleton((uint32_t)ExportID::EpisodeLength); registry.exportColumn( (uint32_t)ExportID::Action); @@ -363,7 +365,7 @@ inline void movementSystem(Engine &e, } else { // Follow expert trajectory const Trajectory &trajectory = e.get(agent_iface.e); - CountT curStepIdx = getCurrentStep(e.get(agent_iface.e)); + CountT curStepIdx = getCurrentStep(e, e.get(agent_iface.e)); position.x = trajectory.positions[curStepIdx].x; position.y = trajectory.positions[curStepIdx].y; position.z = 1; @@ -596,7 +598,7 @@ inline void doneSystem(Engine &ctx, Done &done = ctx.get(agent_iface.e); Info &info = ctx.get(agent_iface.e); int32_t num_remaining = steps_remaining.t; - if (num_remaining == consts::episodeLen && done.v != 1) + if (num_remaining == ctx.singleton().episode_length && done.v != 1) { // Make sure to not reset an agent's done flag done.v = 0; return; @@ -630,7 +632,7 @@ void collisionDetectionSystem(Engine &ctx, // Case: If an expert agent is in an invalid state, we need to ignore the collision detection for it. if (controlledState == false) { - auto currStep = getCurrentStep(ctx.get(agent_iface.value().e)); + auto currStep = getCurrentStep(ctx, ctx.get(agent_iface.value().e)); auto &validState = ctx.get(agent_iface.value().e).valids[currStep]; if (!validState) { @@ -997,6 +999,9 @@ Sim::Sim(Engine &ctx, for (auto i = 0; i < consts::kMaxAgentCount; i++) { deletedAgents.deletedAgents[i] = -1; } + + auto& episodeLength = ctx.singleton(); + episodeLength.episode_length = init.params->episode_length; // Creates agents, walls, etc. createPersistentEntities(ctx); diff --git a/src/sim.hpp b/src/sim.hpp index 16ddc2e81..8bdcad5ef 100755 --- a/src/sim.hpp +++ b/src/sim.hpp @@ -41,6 +41,7 @@ enum class ExportID : uint32_t { DeletedAgents, MapName, ScenarioId, + EpisodeLength, NumExports }; diff --git a/src/types.hpp b/src/types.hpp index 0ac27d418..199452a48 100755 --- a/src/types.hpp +++ b/src/types.hpp @@ -90,6 +90,11 @@ namespace madrona_gpudrive int32_t reset; }; + struct EpisodeLength + { + int32_t episode_length; + }; + struct ResetMap { int32_t reset; }; @@ -337,15 +342,15 @@ namespace madrona_gpudrive struct Trajectory { - madrona::math::Vector2 positions[consts::kTrajectoryLength]; - madrona::math::Vector2 velocities[consts::kTrajectoryLength]; - float headings[consts::kTrajectoryLength]; - float valids[consts::kTrajectoryLength]; - Action inverseActions[consts::kTrajectoryLength]; + madrona::math::Vector2 positions[consts::maxEpisodeLength]; + madrona::math::Vector2 velocities[consts::maxEpisodeLength]; + float headings[consts::maxEpisodeLength]; + float valids[consts::maxEpisodeLength]; + Action inverseActions[consts::maxEpisodeLength]; static inline void zero(Trajectory& traj) { - for (int i = 0; i < consts::kTrajectoryLength; i++) + for (int i = 0; i < consts::maxEpisodeLength; i++) { traj.positions[i] = {0, 0}; traj.velocities[i] = {0, 0}; @@ -356,7 +361,7 @@ namespace madrona_gpudrive } }; - const size_t TrajectoryExportSize = 2 * 2 * consts::kTrajectoryLength + 2 * consts::kTrajectoryLength + ActionExportSize * consts::kTrajectoryLength; + const size_t TrajectoryExportSize = 2 * 2 * consts::maxEpisodeLength + 2 * consts::maxEpisodeLength + ActionExportSize * consts::maxEpisodeLength; static_assert(sizeof(Trajectory) == sizeof(float) * TrajectoryExportSize); diff --git a/src/viewer.cpp b/src/viewer.cpp index 5e25837aa..3a3fcc203 100644 --- a/src/viewer.cpp +++ b/src/viewer.cpp @@ -201,7 +201,7 @@ int main(int argc, char *argv[]) mgr.step(); stepCtr++; - if(stepCtr % consts::episodeLen == 0) { + if(stepCtr % 91 == 0) { // TODO: change to episode length mgr.reset({0}); }