fix CI

waymax/visualization/gokart_viz.py

@@ -16,14 +16,14 @@
 
 from functools import partial
 from typing import Any, Iterable, List, Optional
+import typing
 
 import jax
 import matplotlib
 
 from matplotlib.animation import FuncAnimation
 import numpy as np
 
-from gocarx.env.obs.obs import GokartObs
 from waymax import config as waymax_config
 from waymax import datatypes
 from waymax.datatypes import operations
@@ -34,6 +34,8 @@
 
 import matplotlib.pyplot as plt
 
+GokartObs = typing.TypeVar("GokartObs")
+
 
 def create_video_simulator_state(
     state: datatypes.SimulatorState,