Adding full z-coordinate awareness by Aditya-Gupta26 · Pull Request #161 · Emerge-Lab/PufferDrive

Aditya-Gupta26 · 2025-11-30T04:04:52Z

This PR is aimed at adding z-coordinate support for agents from training to inference.
The agents now are additionally aware of the z-coordinate positions of their surroundings and themselves, which should modify the training behavior. Also, a 2.5D z-collision check support has been added which also looks at the z-coordinate of the agents before marking a collision.

The agents also briefly turn green in situations where they overlap each other (x and y coordinate match).

greptile-apps · 2025-11-30T04:08:35Z

Greptile Overview

Greptile Summary

This PR adds comprehensive z-coordinate awareness to the autonomous driving simulation, enabling agents to understand elevation and implement 2.5D collision detection. The implementation spans from data loading through neural network observations to collision checking.

Key Changes:

Expanded observation space from 7 to 8 features (classic) or 10 to 11 features (jerk) to include z-coordinates for goals, agents, and road segments
Implemented dynamic z-coordinate updates where agents' elevation follows nearby road surfaces by averaging the z-values of nearest road segments
Added 2.5D collision detection: AABB collision in x-y plane followed by z-axis overlap check
Introduced aabb_collision_state to track when agents overlap in x-y but not z (visualized as green)
Updated neural network architectures (PyTorch and C) to process 8-feature partner observations and 15-feature road observations
Added 8 new CARLA town map JSON files and updated binary map format with z-coordinate data
Expanded collision detection grid from 5x5 to 15x15 for better spatial awareness

Critical Issues Found:

Duplicate DepthPoint typedef will cause compilation failure
Uninitialized array elements in road_neighbours lead to undefined behavior in qsort
Multiple debug print statements left in production code

Confidence Score: 2/5

This PR has critical compilation and runtime bugs that must be fixed before merging
Score reflects two critical bugs: duplicate typedef causing compilation failure and uninitialized array causing undefined behavior. While the overall architecture and approach are sound, these issues will prevent the code from compiling or cause unpredictable crashes at runtime.
pufferlib/ocean/drive/drive.h requires immediate attention - contains both the duplicate typedef and the uninitialized array bug. Debug prints in pufferlib/ocean/torch.py, pufferlib/ocean/drive/drive.py, and pufferlib/ocean/drive/visualize.c should also be removed.

Important Files Changed

File Analysis

Filename	Score	Overview
pufferlib/ocean/drive/drive.h	2/5	Core z-coordinate implementation with critical bugs: duplicate typedef definition and uninitialized array leading to undefined behavior
pufferlib/ocean/drive/drive.py	4/5	Updated observation space dimensions from 7 to 8 features, added debug print statements that should be removed
pufferlib/ocean/torch.py	4/5	Neural network updated for 8-feature observations and 15-feature road encoding, includes debug prints that should be removed
pufferlib/ocean/drive/drivenet.h	5/5	C neural network updated for 8-feature observations, correctly synchronized with PyTorch implementation
pufferlib/ocean/drive/visualize.c	4/5	Added debug print statements and fflush calls, one incorrect printf statement

Sequence Diagram

sequenceDiagram
    participant Env as Drive Environment
    participant Agent as Agent Entity
    participant Road as Road Segments
    participant Collision as Collision System
    participant NN as Neural Network

    Note over Env: Initialize with z-coordinates
    Env->>Agent: Load trajectory data (x, y, z)
    Env->>Road: Load road geometry (x, y, z)
    Env->>Env: Compute world_mean_z
    Env->>Agent: Normalize z -= world_mean_z
    Env->>Road: Normalize z -= world_mean_z
    
    Note over Env: Step Loop
    Agent->>NN: Compute observations (8 features: goal_x, goal_y, goal_z, ...)
    NN->>Agent: Return actions
    Agent->>Agent: Update position (x, y)
    
    Note over Agent: Z-Coordinate Update
    Agent->>Road: Find nearby road segments (checkNeighbors)
    Road->>Agent: Return list of road entities
    Agent->>Agent: Compute z-distance to each road (compute_z_distance_to_road_segment)
    Agent->>Agent: Sort by distance (qsort)
    Agent->>Agent: Average z from nearest segments
    Agent->>Agent: Update agent->z
    
    Note over Collision: 2.5D Collision Check
    Agent->>Collision: Check collision with other agents
    Collision->>Collision: check_aabb_collision (x, y planes)
    alt AABB collision detected
        Collision->>Agent: Set aabb_collision_state = 1
        Collision->>Collision: check_z_collision (z axis overlap)
        alt Z overlap detected
            Collision->>Agent: Set collision_state = 1
        else No Z overlap
            Collision->>Agent: No collision (agents on different elevations)
        end
    end
    
    Agent->>Env: Return observations, rewards, done

greptile-apps

_{15 files reviewed, 5 comments}

_{Edit Code Review Agent Settings | Greptile}

greptile-apps · 2025-11-30T04:08:31Z

pufferlib/ocean/drive/drive.h

            a_long_new = clip(a_long_new, -5.0f, 2.5f);
        }

+        typedef struct {
+            float dis;


syntax: Duplicate DepthPoint typedef - already defined at line ~424. This will cause compilation errors.

Suggested change

a_long_new = clip(a_long_new, -5.0f, 2.5f);

}

typedef struct {

float dis;

Prompt To Fix With AI

This is a comment left during a code review. Path: pufferlib/ocean/drive/drive.h Line: 1828:1832 Comment: **syntax:** Duplicate `DepthPoint` typedef - already defined at line ~424. This will cause compilation errors. ```suggestion ``` How can I resolve this? If you propose a fix, please make it concise.

greptile-apps · 2025-11-30T04:08:32Z

pufferlib/ocean/torch.py

+        road_onehot = F.one_hot(road_categorical.long(), num_classes=8)  # Shape: [batch, 200, 8]
+        print(road_onehot.shape)
+        print(road_continuous.shape)


style: Debug print statements left in production code

Suggested change

road_onehot = F.one_hot(road_categorical.long(), num_classes=8) # Shape: [batch, 200, 8]

print(road_onehot.shape)

print(road_continuous.shape)

Prompt To Fix With AI

This is a comment left during a code review. Path: pufferlib/ocean/torch.py Line: 78:80 Comment: **style:** Debug print statements left in production code ```suggestion ``` How can I resolve this? If you propose a fix, please make it concise.

greptile-apps · 2025-11-30T04:08:33Z

pufferlib/ocean/drive/drive.py

+        print(len(map_data.get("objects", [])))
+        print(len(map_data.get("roads", [])))


style: Debug print statements left in production code

Suggested change

print(len(map_data.get("objects", [])))

print(len(map_data.get("roads", [])))

Prompt To Fix With AI

This is a comment left during a code review. Path: pufferlib/ocean/drive/drive.py Line: 425:426 Comment: **style:** Debug print statements left in production code ```suggestion ``` How can I resolve this? If you propose a fix, please make it concise.

greptile-apps · 2025-11-30T04:08:34Z

pufferlib/ocean/drive/visualize.c


    if (render_topdown) {
        printf("Recording topdown view...\n");
+        fflush(stdout);


style: Debug printf left in code - env.observations is a pointer, this will print the address not the size

Prompt To Fix With AI

This is a comment left during a code review. Path: pufferlib/ocean/drive/visualize.c Line: 361:361 Comment: **style:** Debug printf left in code - `env.observations` is a pointer, this will print the address not the size How can I resolve this? If you propose a fix, please make it concise.

greptile-apps · 2025-11-30T04:08:35Z

pufferlib/ocean/drive/drive.h

+    //To update agent's z-coordinate based on road elevation of 5 nearest elements
+    GridMapEntity entity_list[MAX_ENTITIES_PER_CELL*100];  // Array big enough for all neighboring cells
+    int list_size = checkNeighbors(env, agent->x, agent->y, entity_list, MAX_ENTITIES_PER_CELL*100, collision_offsets, 100);
+    DepthPoint road_neighbours[list_size];
+    int max_check = (list_size < 10) ? list_size : 10;
+    int diffarray[max_check-1];
+    if(list_size>0){
+        //store an array masuring the distance of the agent with each road segment nearby
+        for(int i = 0;i<list_size;i++){
+            if(entity_list[i].entity_idx == -1 ) continue;
+            Entity* entity = &env->entities[entity_list[i].entity_idx];
+            if(entity->type == ROAD_EDGE || entity->type == ROAD_LANE || entity->type == ROAD_LINE){


logic: road_neighbours array not fully initialized. Loop only sets elements where entity->type is ROAD_EDGE/ROAD_LANE/ROAD_LINE, leaving other indices uninitialized. When qsort reads these, it accesses undefined values.

// Initialize array first DepthPoint road_neighbours[list_size]; memset(road_neighbours, 0, list_size * sizeof(DepthPoint)); int valid_count = 0; for(int i = 0; i < list_size; i++){ if(entity_list[i].entity_idx == -1) continue; Entity* entity = &env->entities[entity_list[i].entity_idx]; if(entity->type == ROAD_EDGE || entity->type == ROAD_LANE || entity->type == ROAD_LINE){ road_neighbours[valid_count++] = compute_z_distance_to_road_segment(agent, entity); } } qsort(road_neighbours, valid_count, sizeof(DepthPoint), compare_depthpoint);

Prompt To Fix With AI

This is a comment left during a code review. Path: pufferlib/ocean/drive/drive.h Line: 1685:1696 Comment: **logic:** `road_neighbours` array not fully initialized. Loop only sets elements where `entity->type` is ROAD_EDGE/ROAD_LANE/ROAD_LINE, leaving other indices uninitialized. When `qsort` reads these, it accesses undefined values. ``` // Initialize array first DepthPoint road_neighbours[list_size]; memset(road_neighbours, 0, list_size * sizeof(DepthPoint)); int valid_count = 0; for(int i = 0; i < list_size; i++){ if(entity_list[i].entity_idx == -1) continue; Entity* entity = &env->entities[entity_list[i].entity_idx]; if(entity->type == ROAD_EDGE || entity->type == ROAD_LANE || entity->type == ROAD_LINE){ road_neighbours[valid_count++] = compute_z_distance_to_road_segment(agent, entity); } } qsort(road_neighbours, valid_count, sizeof(DepthPoint), compare_depthpoint); ``` How can I resolve this? If you propose a fix, please make it concise.

daphne-cornelisse · 2025-12-01T17:56:40Z

pufferlib/ocean/drive/drive.c

    float* observations = calloc(num_agents*num_obs, sizeof(float));
    for (int i = 0; i < num_obs*num_agents; i++) {
-        observations[i] = i % 7;
+        observations[i] = i % 8;


We discussed this on Slack, but after rebasing, you can simply increment these

// Ego features depend on dynamics model #define EGO_FEATURES_CLASSIC 7 #define EGO_FEATURES_JERK 10

daphne-cornelisse · 2025-12-01T17:57:38Z

pufferlib/ocean/drive/drive.h

-    {-2,  0}, {-1,  0}, {0,  0}, {1,  0}, {2,  0},  // Middle row (including center)
-    {-2,  1}, {-1,  1}, {0,  1}, {1,  1}, {2,  1},  // Fourth row
-    {-2,  2}, {-1,  2}, {0,  2}, {1,  2}, {2,  2}   // Bottom row
+static const int collision_offsets[225][2] = {


Wow, why is this necessary? Can this be done more simply?

I wanted to increase the size of the grid so at all points, there are always enough road lanes/lines/edges which the agent can find. With the older (smaller) grid, I noticed that it was unable to find enough agents, and would update the z value based on 1-2 surrounding points, which might be inaccurate.

daphne-cornelisse · 2025-12-01T17:58:15Z

pufferlib/ocean/drive/drive.h

+    float init_goal_z;
    int mark_as_expert;
    int collision_state;
+    int aabb_collision_state;


Why does aabb_collision_state need to be added to the timespec struct?

aabb_collision_state is meant to be set only when the x and y coordinates of agents overlap (2D collision). I added this specifically so that the agent knows when to turn green.

daphne-cornelisse · 2025-12-01T17:58:51Z

pufferlib/ocean/drive/drive.h

        fread(&entities[i].length, sizeof(float), 1, file);
        fread(&entities[i].height, sizeof(float), 1, file);
        fread(&entities[i].goal_position_x, sizeof(float), 1, file);
+        //fread(&entities[i].goal_position_z, sizeof(float), 1, file);


Is there a goal position z?

Yes, support for goal_position_z exists, I also added support to update it in case of reset, or while computing new goals

daphne-cornelisse · 2025-12-01T18:01:11Z

pufferlib/ocean/drive/drive.h

-            car_collided_with_index = index;
-            break;
+            agent->aabb_collision_state = 1;
+            if(check_z_collision(agent,entity)){


Why not include the z-axis into the check_aabb_collision? What is the rationale behind creating a separate function for the z-axis?

So the way collision check works now is that we first check for a 2D collision (x and y values overlap). Only if there is a 2D collision, we go ahead and check if the z-values of the agents are close enough for a collision (the check_z_collision function).
We can include the z-axis check inside check_aabb_collision, I just thought separating this would be cleaner.

…ring mode

…ection

greptile-apps bot reviewed Nov 30, 2025

View reviewed changes

daphne-cornelisse self-requested a review December 1, 2025 17:55

daphne-cornelisse reviewed Dec 1, 2025

View reviewed changes

Aditya-Gupta26 and others added 24 commits December 20, 2025 23:18

Updates after rebasing

eee8094

Updates for z coordinate support

2e57765

Adding full z-coordinate awareness

6357b8f

Removing unneccsary print statements

0146653

Adding full z-coordinate awareness

8c41106

Adding full z-coordinate awareness

dadc8d0

Resetting visualize.c to fix render issues

cb65965

Removing testing override

a33eead

Removing testing override

1ba8911

Removing testing override

84c5358

Removing testing override

531cc01

Removing testing override

ed9f7cd

Removing testing override

b73d8d7

Removing testing override

95aac46

Remove files that should not be included

79dd064

Fixing some comments and print statements for debugging

2b54ff1

Removing unwanted comments and files

d3df1e5

Testing with 3D carla

22a77a0

Adding agent initialzed maps for training

5df5ca3

Adding controllable agents to maps

2fe290f

Debugging

2047237

Fixes

47d1d5f

Fixing offroad logic

702c320

Fixing duplicates

0aea6bd

Updating visualize to match latest main

262bfa7

Aditya-Gupta26 force-pushed the aditya/z_coord_carla branch from 98bf017 to 262bfa7 Compare December 22, 2025 03:40

Aditya-Gupta26 added 5 commits December 22, 2025 00:11

Fixing drivenet

f2263d4

Adding weights file

ca53978

Adding weights file

e50756d

All merged changes and fixes

d4c1a50

Adding latest 3D CARLA maps for testing

85ec1ec

daphne-cornelisse added this to the 3.0 milestone Dec 25, 2025

Aditya-Gupta26 and others added 12 commits January 6, 2026 17:54

Fixing one hot error

73559df

Minor

fc99f04

Removing print statements

fdb08b8

Updating over latest main branch changes and removing cinematic rende…

0dcb5c2

…ring mode

Adding dynamic grid entities changes

04fd8d0

Minor fix

63937e2

Adding back dynamic grid entities

e9b2b06

Adding latest changes from 3.0 which makes codebase more flexible

73d679e

Adding memory leak fix

c78b055

Fix edge-case with no entity around

4cd0c6d

Test out with 5*5 grid and dynamic grid

7f91f35

Adding a different offset for z calculation and collision/offroad det…

4b09a4b

…ection

Aditya-Gupta26 changed the base branch from main to 3.0 January 13, 2026 21:41

Aditya-Gupta26 added 3 commits January 13, 2026 17:12

Fixing map observation rendering

6a27030

Rendering bug fix

f4c7e4c

Updating 3D carla maps with latest .json changes

6991440

	road_onehot = F.one_hot(road_categorical.long(), num_classes=8) # Shape: [batch, 200, 8]
	print(road_onehot.shape)
	print(road_continuous.shape)

		print(len(map_data.get("objects", [])))
		print(len(map_data.get("roads", [])))

Conversation

Aditya-Gupta26 commented Nov 30, 2025

Uh oh!

greptile-apps bot commented Nov 30, 2025

Greptile Overview

Greptile Summary

Confidence Score: 2/5

Important Files Changed

Sequence Diagram

Uh oh!

greptile-apps bot left a comment

Choose a reason for hiding this comment

Uh oh!

greptile-apps bot Nov 30, 2025

Choose a reason for hiding this comment

Uh oh!

greptile-apps bot Nov 30, 2025

Choose a reason for hiding this comment

Uh oh!

greptile-apps bot Nov 30, 2025

Choose a reason for hiding this comment

Uh oh!

greptile-apps bot Nov 30, 2025

Choose a reason for hiding this comment

Uh oh!

greptile-apps bot Nov 30, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants