Speed limit and lane center/angle observation

[Aditya-Gupta26]

This PR adds the following -

• Speed limit infra (for new data format use later) and currently uses a hardcoded speed limit value.
• Adds lane center/angle information by computing the closest lane segment and a weighted average for smooth angle estimation

These are added onto the observation vector for both JERK and CLASSIC dynamic models

(This is derived from Turbocharge PR)

Note : This PR is not meant to add reward conditioning for lane/speed alignment as of yet. That will be catered to in a separate PR

[greptile-apps]

Greptile Overview

Greptile Summary

This PR extends the ego observation for both dynamics models by adding (1) a speed-limit feature (currently hardcoded) and (2) lane-center distance + lane-angle alignment computed from the closest lane segment with a small multi-segment heading smoothing. The C simulator (drive.h) now computes these lane metrics during compute_agent_metrics and appends them in compute_observations, while the PyTorch policy (torch.py) switches to using env.ego_features so the encoder input size matches the environment’s ego feature count.

Key integration risk: metrics_array and current_lane_index semantics are being expanded/repurposed, but the reset/respawn/initialization paths and the existing “closest lane” logic are not consistently updated, which can yield stale/incorrect lane metrics in observations and inconsistent lane tracking across timesteps.

Confidence Score: 2/5

• Not safe to merge until lane metrics state/reset issues are fixed
• The observation shape changes are mostly wired through (env.ego_features in torch.py and updated EGO_FEATURES_* constants), but drive.h introduces definite correctness issues: lane selection state is overwritten by existing closest-lane logic, and new lane metric slots aren’t consistently reset/initialized on step reset, set_start_position, and respawn. These lead to inconsistent lane tracking and stale/incorrect observations.
• pufferlib/ocean/drive/drive.h (metrics reset/init and lane index semantics); pufferlib/ocean/drive/datatypes.h (metrics_array documentation alignment)

Important Files Changed

Filename	Overview
pufferlib/ocean/drive/datatypes.h	Expanded Agent.metrics_array to 11; comment still describes 10 entries, requiring alignment with new metric indices and reset paths.
pufferlib/ocean/drive/drive.h	Adds lane segment selection + lane center/angle observation fields and expands ego feature dims; introduces several state/metrics reset and lane-index consistency bugs.
pufferlib/ocean/torch.py	Uses env.ego_features for ego encoder input size, matching updated observation ego dimension constants.

Sequence Diagram

sequenceDiagram
    participant PyEnv as drive.py (Gym env)
    participant CEnv as drive.h (C sim)
    participant Agent as Agent/metrics_array
    participant Policy as torch.py (NN policy)

    PyEnv->>CEnv: allocate() uses EGO_FEATURES_* to size observations
    loop each step
        CEnv->>CEnv: compute_agent_metrics(agent)
        CEnv->>Agent: write metrics_array[LANE_DIST_IDX], [LANE_ANGLE_IDX]
        CEnv->>CEnv: compute_observations()
        CEnv->>CEnv: obs[ego] includes speed_limit, lane_center_dist, lane_angle
        CEnv-->>PyEnv: observations buffer
        PyEnv->>Policy: forward(observations)
        Policy->>Policy: slice using env.ego_features
    end

Loading

[greptile-apps]

_{3 files reviewed, 5 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Metrics indices now mismatch

metrics_array is resized to 11 in datatypes.h, but the new indices in drive.h imply you’re now writing up to index 10 (LANE_ALIGNED_IDX). Existing comments in Agent still describe only 10 items, and several reset/respawn paths only clear a subset of indices. This will leave stale values (e.g., LANE_DIST_IDX, LANE_ANGLE_IDX) across steps/episodes and makes the meaning of each index ambiguous. Please update the Agent comment to list all 11 entries and ensure all codepaths that reset metrics zero the full [0..LANE_ALIGNED_IDX] range (or explicitly initialize new slots).

[greptile-apps]

Additional Comments (4)

pufferlib/ocean/drive/drive.h
New metrics not reset

reset_agent_metrics only clears COLLISION_IDX, OFFROAD_IDX, and LANE_ALIGNED_IDX. With the new lane center/angle observation derived from metrics_array[LANE_DIST_IDX]/[LANE_ANGLE_IDX], those values can persist from prior timesteps/episodes for agents that don’t get a lane assignment this step, leading to incorrect observations. Consider zeroing the full metrics array (or at least indices 0..10) here.

---

pufferlib/ocean/drive/drive.h
Lane choice overwritten later

The new “best candidate” lane selection updates agent->current_lane_index/current_lane_geometry_idx (and writes LANE_DIST_IDX/LANE_ANGLE_IDX), but immediately after that the existing block sets current_lane_index again based on the closest ROAD_LANE segment (closest_lane_entity_idx). This means current_lane_index will not consistently correspond to the lane used for the center/angle metrics, and the hysteresis logic (which depends on current_lane_index) won’t behave as intended across steps. Consider separating the concepts (e.g., keep a dedicated “current drivable lane” index for the new metrics) or stop overwriting current_lane_lane_index after selecting the best candidate.

---

pufferlib/ocean/drive/drive.h
Agent init misses new slots

In set_start_position, only a few metrics_array indices are initialized (COLLISION, OFFROAD, REACHED_GOAL, LANE_ALIGNED). With the expanded metrics_array and new observation fields depending on LANE_DIST_IDX/LANE_ANGLE_IDX, those entries start uninitialized/garbage for freshly initialized agents until compute_agent_metrics runs. Consider zeroing the entire metrics_array here (or explicitly initializing indices 0..10).

---

pufferlib/ocean/drive/drive.h
Respawn leaves lane metrics stale

respawn_agent clears only a few metrics_array indices, but the new observation uses LANE_DIST_IDX/LANE_ANGLE_IDX. After respawn, those entries can retain pre-respawn values until the next compute_agent_metrics update, producing incorrect lane center/angle observations right after respawn. Consider clearing the full metrics_array (or at least the new indices) during respawn.