Z-axis support integration

[Aditya-Gupta26]

This PR contains the following new features -

• Z-coordinate awareness (Agents have a 3D spatial view around them now)
  • Updated the number of observations fed to the network, which now includes z positions of goals, agents, map points
• Updating Z-coordinate values for agents
  • Agents will have a z-coordinate value measured as the average of a certain number of nearest points within a grid
• 2.5D collision/offroad checking
  • Z-level of agents will be checked before flagging agents as collided/offroad
• 3D Rendering support for visualization
  • The demo() function in drive.c, and the other render views can now support 3D contour depending on the map

Some Carla maps are also added/updated to match their latest versions.

[Aditya-Gupta26]

Screen.Recording.2026-01-16.at.3.40.01.PM.mov

Here's how everything looks (Its with goal resampling).

[greptile-apps]

Greptile Summary

This PR adds comprehensive 3D spatial awareness to the driving environment, enabling agents to navigate maps with elevation changes like bridges and overpasses.

Key Changes:

• 3D Observations: Agents now receive z-coordinates for ego state, partner vehicles, and road segments. Feature counts increased from 7→8 (classic) and 10→11 (jerk dynamics) to include z-position.
• Z-Level Computation: Agent elevation is dynamically computed by averaging the z-coordinates of the 20 nearest road segments, with outlier detection to handle elevation transitions.
• 2.5D Collision Detection: Two-stage collision system first checks 2D AABB overlap, then validates z-axis overlap. Includes 4m buffer for offroad detection to prevent false positives on bridges.
• 3D Rendering: Visualization updated to render contoured maps with proper camera positioning based on agent z-coordinates.
• Map Data: Updated CARLA maps (Town03-07) with 3D geometry and added new map processing pipeline.

Critical Issues:

• Uninitialized memory bug in z-level computation: The road_neighbours array is partially filled but fully sorted, causing undefined behavior.
• Potential edge case when no valid road segments are found nearby - agent z-coordinate behavior is unclear.

Other Observations:

• Velocity calculation bug fixed in trajectory generation (was using wrong timestep).
• Network architecture correctly updated in both C and Python to match new observation dimensions.
• Configuration changes align with new CARLA 3D map dataset.

Confidence Score: 2/5

• Critical memory safety bugs in z-coordinate computation require fixing before merge
• Two critical bugs found in drive.h: (1) uninitialized array elements sorted alongside valid data leading to undefined behavior, and (2) potential division by zero edge cases. These are in core physics simulation code that runs every timestep and could cause crashes or incorrect agent behavior. The architectural changes are sound and network updates are correct, but the memory safety issues must be resolved.
• Pay close attention to pufferlib/ocean/drive/drive.h lines 1980-2020 (z-coordinate computation)

Important Files Changed

Filename	Overview
pufferlib/ocean/drive/drive.h	Core 3D spatial logic added including z-coordinate computation, 2.5D collision detection, and observation updates. Contains critical bugs in z-level calculation and potential uninitialized memory access.
pufferlib/ocean/drive/drivenet.h	Network architecture correctly updated to handle 8 road features (was 7) and proper one-hot encoding indexing for 3D observations.
pufferlib/ocean/torch.py	Correctly updates ego dimension from 7→8 for classic and 10→11 for jerk dynamics to accommodate z-coordinate in observations.
data_utils/carla/generate_carla_agents.py	Fixes velocity calculation bug, improves initial heading computation, and adds debug logging. Velocity fix is important but may affect trajectory quality.

Sequence Diagram

sequenceDiagram
    participant Agent
    participant move_dynamics
    participant checkNeighbors
    participant compute_z_distance
    participant qsort
    participant compute_observations
    participant collision_check
    participant check_z_collision
    
    Note over Agent,check_z_collision: 3D Spatial Processing Flow
    
    Agent->>move_dynamics: Update position (x, y)
    move_dynamics->>checkNeighbors: Get nearby road segments (15x15 grid)
    checkNeighbors-->>move_dynamics: GridMapEntity list
    
    loop For each road segment
        move_dynamics->>compute_z_distance: Calculate z-distance
        compute_z_distance-->>move_dynamics: DepthPoint{dis, z}
    end
    
    move_dynamics->>qsort: Sort by distance
    qsort-->>move_dynamics: Sorted DepthPoints
    
    Note over move_dynamics: Find max gap in distances<br/>Average z-coords up to gap
    
    move_dynamics->>Agent: Set agent.z
    
    Agent->>compute_observations: Generate observations
    
    Note over compute_observations: Ego: 8 features (x,y,z,vx,vy,heading_x,heading_y,speed)<br/>Partners: 8 features each (rel_x,rel_y,rel_z,w,l,heading_x,heading_y,speed)<br/>Road: 8 features each (rel_x,rel_y,rel_z,length,width,cos,sin,type)
    
    compute_observations-->>Agent: 3D observations
    
    Agent->>collision_check: Check collisions
    collision_check->>collision_check: AABB collision (2D)
    
    alt AABB collision detected
        collision_check->>Agent: Set aabb_collision_state=1
        collision_check->>check_z_collision: Check z-overlap
        
        alt Z-ranges overlap
            check_z_collision-->>collision_check: True collision
            collision_check->>Agent: Set collision_state=1
        else Z-ranges separated
            check_z_collision-->>collision_check: No collision
            Note over Agent: Different vertical levels<br/>(e.g., bridge/underpass)
        end
    end

Loading

[greptile-apps]

Additional Comments (1)

1. pufferlib/ocean/drive/drive.h, line 93-95 (link)

   style: Feature counts updated from 7→8 for classic and 10→11 for jerk to include z-coordinate in observations.

_{20 files reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}