`[:, 2]` indexing in `convert_absolute_to_relative_se2_array` corrupts batched `(M, T, 3)` input and skips heading normalization

[tjyuyao]

Summary

navsim/planning/simulation/planner/pdm_planner/utils/pdm_geometry_utils.py:95:

points_rel[:, 2] = normalize_angle(points_rel[:, 2])

uses [:, 2] instead of [..., 2]. For batched input of shape (M, T, 3) (M trajectories × T timesteps), points_rel[:, 2] selects timestep index 2's whole row (M, 3) rather than the heading column (M, T). Two consequences:

1. normalize_angle is elementwise (np.arctan2(np.sin(x), np.cos(x))), so timestep 2's x/y coordinates are also "angle-normalized" into (-π, π], corrupting positions (e.g. x=30 m → ≈4.87).
2. Heading normalization never runs for the remaining T−1 timesteps — the intended behavior of this line is silently skipped.

Minimal repro

import numpy as np
from nuplan.common.actor_state.state_representation import StateSE2
from navsim.planning.simulation.planner.pdm_planner.utils.pdm_geometry_utils import (
    convert_absolute_to_relative_se2_array,
)

origin = StateSE2(0.0, 0.0, 0.0)
states = np.zeros((3, 4, 3))
states[..., 0] = 30.0  # x = 30 m for all trajectories and timesteps

out = convert_absolute_to_relative_se2_array(origin, states)
print(out[:, 2, 0])   # corrupted: ≈4.87 (atan2(sin 30, cos 30)), expected 30.0
print(out[0, :, 2])   # heading untouched for all timesteps except index 2

Suggested fix

points_rel[..., 2] = normalize_angle(points_rel[..., 2])

Backward compatible: ... degenerates to : for the existing 2-D (N, 3) callers.

Impact

In-repo callers (all 4 in navsim/common/dataclasses.py) pass 2-D (N, 3) single-trajectory arrays, so the bug is latent — current official code paths are unaffected. However, the function's contract (docstring: "array of SE2 states with (x,y,θ) in last dim"; assert only checks shape[-1] == 3) places no dimension limit, so batched (M, T, 3) input is within the documented contract and silently corrupts timestep-2 positions while skipping heading normalization for all other timesteps — a silent correctness hazard for any future batched usage (e.g. converting a batch of proposal trajectories in a scoring loop), easy to miss when downstream code does not consume timestep 2.