From 133eefe89b3b9e6d4ba5ecb688bd84c045721a18 Mon Sep 17 00:00:00 2001 From: SS47816 Date: Wed, 31 Jan 2024 19:51:32 +0800 Subject: [PATCH] fix: array index when center_agent_idx is not sdc --- waymax/visualization/viz.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/waymax/visualization/viz.py b/waymax/visualization/viz.py index edc6871..c567aab 100644 --- a/waymax/visualization/viz.py +++ b/waymax/visualization/viz.py @@ -282,9 +282,10 @@ def plot_simulator_state( current_xy = traj.xy[:, state.timestep, :] if viz_config.center_agent_idx == -1: xy = current_xy[state.object_metadata.is_sdc] + origin_x, origin_y = xy[0, :2] else: xy = current_xy[viz_config.center_agent_idx] - origin_x, origin_y = xy[0, :2] + origin_x, origin_y = xy[:2] ax.axis(( origin_x - viz_config.back_x, origin_x + viz_config.front_x,