Fix bug in agent view for jerk dynamics model

[daphne-cornelisse]

No description provided.

[greptile-apps]

Greptile Summary

Fixes observation buffer indexing in the draw_agent_obs visualization function to correctly handle the jerk dynamics model, which uses 10 ego features instead of the classic model's 7.

• Replaced hardcoded 7 with PARTNER_FEATURES constant when advancing through partner observations
• Replaced hardcoded 7 + 7 * (MAX_AGENTS - 1) with ego_dim + PARTNER_FEATURES * (MAX_AGENTS - 1) for calculating the map observations start index
• One remaining hardcoded 7 at line 2527 was missed and should also be updated for complete consistency

Confidence Score: 4/5

• This PR is safe to merge but has one missed instance of the same bug pattern that should be fixed.
• The fix correctly addresses the core issue with jerk dynamics model observation indexing. However, one hardcoded 7 at line 2527 in the same function was missed, which could cause the same visualization bug when skipping null partner observations in jerk mode.
• Line 2527 in pufferlib/ocean/drive/drive.h needs the same fix applied for complete consistency.

Important Files Changed

Filename	Overview
pufferlib/ocean/drive/drive.h	Fixes hardcoded observation indexing in draw_agent_obs to use ego_dim and PARTNER_FEATURES constants for proper jerk dynamics model support. One instance at line 2527 still uses hardcoded 7 instead of PARTNER_FEATURES.

Sequence Diagram

sequenceDiagram
    participant User as draw_agent_obs()
    participant Ego as Ego Observations
    participant Partner as Partner Observations
    participant Map as Map Observations
    
    Note over User: Visualization function for agent observations
    
    User->>Ego: Read ego observations (0 to ego_dim)
    Note over Ego: Classic: 7 features<br/>Jerk: 10 features
    
    User->>Partner: Iterate partner observations
    loop For each partner (MAX_AGENTS - 1)
        Partner-->>User: Read PARTNER_FEATURES (7) values
        Note over Partner: obs_idx += PARTNER_FEATURES
    end
    
    User->>Map: Calculate map_start_idx
    Note over Map: map_start_idx = ego_dim +<br/>PARTNER_FEATURES * (MAX_AGENTS - 1)
    
    loop For each road segment
        Map-->>User: Read 7 road feature values
    end

Loading

[greptile-apps]

Additional Comments (1)

1. pufferlib/ocean/drive/drive.h, line 2527 (link)

   logic: This hardcoded 7 should also be PARTNER_FEATURES for consistency with the fix at line 2633. Both are in the same loop advancing obs_idx through partner observations.

_{1 file reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}