Skip to content

[WIP] Add support for dynamic episode lengths #407

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[WIP] Add support for dynamic episode lengths #407

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions src/bindings.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Copy link
Contributor

@daphne-cornelisse daphne-cornelisse Apr 23, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can't we stick to the original naming of episodeLen? The fact that an agent can be terminated before the episodeLen seems implied

m.attr("numLidarSamples") = consts::numLidarSamples;
m.attr("vehicleScale") = consts::vehicleLengthScale;

Expand Down Expand Up @@ -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);
Copy link
Contributor

@daphne-cornelisse daphne-cornelisse Apr 23, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We already have the steps remaining tensor, is the only diff that this one is counting up?

}

}
4 changes: 2 additions & 2 deletions src/consts.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Copy link
Contributor

@daphne-cornelisse daphne-cornelisse Apr 23, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we should stick to episodeLen

Also, let's keep the default value


// Number of lidar samples, arranged in circle around agent
inline constexpr madrona::CountT numLidarSamples = 50;
Expand All @@ -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.
Copy link
Contributor

@daphne-cornelisse daphne-cornelisse Apr 23, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Even if we support variable episode lengths, the log replay trajectories in the scenes will remain of shape 91.

Why is this deleted? There is an important difference between the length of the episode and the length of the data in the simulator


inline constexpr madrona::CountT kMaxRoadGeometryLength = 1810;

Expand Down
3 changes: 2 additions & 1 deletion src/init.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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;
Copy link
Contributor

@daphne-cornelisse daphne-cornelisse Apr 23, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Seeming inconsistency here: episode_length = 91 here and max episode length is 200?

};

struct WorldInit
Expand Down
9 changes: 5 additions & 4 deletions src/json_serialization.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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;
Expand All @@ -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;
Expand All @@ -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)
Copy link
Contributor

@daphne-cornelisse daphne-cornelisse Apr 23, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What does this block of code do?

{
v.get_to(obj.valid[i]);
++i;
Expand Down
4 changes: 2 additions & 2 deletions src/level_gen.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ static void registerRigidBodyEntity(
ctx.get<broadphase::LeafID>(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<StepsRemaining>(agent_iface).t = steps_remaining;
ctx.get<Done>(agent_iface).v = done;
ctx.get<Reward>(agent_iface).v = 0;
Expand All @@ -46,7 +46,7 @@ static inline void resetAgent(Engine &ctx, Entity agent) {
}
ctx.get<Action>(agent_iface) = getZeroAction(ctx.data().params.dynamicsModel);

resetAgentInterface(ctx, agent_iface, ctx.get<EntityType>(agent), ctx.get<ResponseType>(agent));
resetAgentInterface(ctx, agent_iface, ctx.get<EntityType>(agent), ctx.get<ResponseType>(agent), ctx.singleton<EpisodeLength>().episode_length);

#ifndef GPUDRIVE_DISABLE_NARROW_PHASE
ctx.get<CollisionDetectionEvent>(agent).hasCollided.store_release(0);
Expand Down
6 changes: 6 additions & 0 deletions src/mgr.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
1 change: 1 addition & 0 deletions src/mgr.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
15 changes: 10 additions & 5 deletions src/sim.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<EpisodeLength>().episode_length - stepsRemaining.t;
}

// Register all the ECS components and archetypes that will be
Expand Down Expand Up @@ -68,6 +68,7 @@ void Sim::registerTypes(ECSRegistry &registry, const Config &cfg)
registry.registerSingleton<DeletedAgents>();
registry.registerSingleton<MapName>();
registry.registerSingleton<ScenarioId>();
registry.registerSingleton<EpisodeLength>();

registry.registerArchetype<Agent>();
registry.registerArchetype<PhysicsEntity>();
Expand All @@ -83,6 +84,7 @@ void Sim::registerTypes(ECSRegistry &registry, const Config &cfg)
registry.exportSingleton<DeletedAgents>((uint32_t)ExportID::DeletedAgents);
registry.exportSingleton<MapName>((uint32_t)ExportID::MapName);
registry.exportSingleton<ScenarioId>((uint32_t)ExportID::ScenarioId);
registry.exportSingleton<EpisodeLength>((uint32_t)ExportID::EpisodeLength);

registry.exportColumn<AgentInterface, Action>(
(uint32_t)ExportID::Action);
Expand Down Expand Up @@ -363,7 +365,7 @@ inline void movementSystem(Engine &e,
} else {
// Follow expert trajectory
const Trajectory &trajectory = e.get<Trajectory>(agent_iface.e);
CountT curStepIdx = getCurrentStep(e.get<StepsRemaining>(agent_iface.e));
CountT curStepIdx = getCurrentStep(e, e.get<StepsRemaining>(agent_iface.e));
position.x = trajectory.positions[curStepIdx].x;
position.y = trajectory.positions[curStepIdx].y;
position.z = 1;
Expand Down Expand Up @@ -596,7 +598,7 @@ inline void doneSystem(Engine &ctx,
Done &done = ctx.get<Done>(agent_iface.e);
Info &info = ctx.get<Info>(agent_iface.e);
int32_t num_remaining = steps_remaining.t;
if (num_remaining == consts::episodeLen && done.v != 1)
if (num_remaining == ctx.singleton<EpisodeLength>().episode_length && done.v != 1)
{ // Make sure to not reset an agent's done flag
done.v = 0;
return;
Expand Down Expand Up @@ -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<StepsRemaining>(agent_iface.value().e));
auto currStep = getCurrentStep(ctx, ctx.get<StepsRemaining>(agent_iface.value().e));
auto &validState = ctx.get<Trajectory>(agent_iface.value().e).valids[currStep];
if (!validState)
{
Expand Down Expand Up @@ -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>();
episodeLength.episode_length = init.params->episode_length;
// Creates agents, walls, etc.
createPersistentEntities(ctx);

Expand Down
1 change: 1 addition & 0 deletions src/sim.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ enum class ExportID : uint32_t {
DeletedAgents,
MapName,
ScenarioId,
EpisodeLength,
NumExports
};

Expand Down
19 changes: 12 additions & 7 deletions src/types.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,11 @@ namespace madrona_gpudrive
int32_t reset;
};

struct EpisodeLength
{
int32_t episode_length;
};

struct ResetMap {
int32_t reset;
};
Expand Down Expand Up @@ -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};
Expand All @@ -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;
Copy link
Contributor

@daphne-cornelisse daphne-cornelisse Apr 23, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks incorrect. Even if you increase the episode length, the data shape will always remain the same


static_assert(sizeof(Trajectory) == sizeof(float) * TrajectoryExportSize);

Expand Down
2 changes: 1 addition & 1 deletion src/viewer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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});
}

Expand Down
Loading