How to accelerate the plot process?

[gaoyinfeng]

I successfully set up the environment following the instructions and ran data_demo in jupyter notebook. One thing that bothers me is the time-consuming plotting function, it took a maximum of 9 seconds to plot 1 step, which is unacceptable (although it becomes faster as the episode grows, about 0.2s for 1 step near the end of the episode). I also noticed that the state update of the environment is not as efficient as I expected (about 0.02s to update 1 step), finally it took me 70~80 seconds to make a video which shows a 9 seconds long episode. Maybe it is because something went wrong with my deployment?
Here are my codes and results, I additionally added memory_growth to prevent OOM errors:

BTW, I work on a machine which has an INTEL i7-7800X CPU and NVIDIA RTX 2080Ti GPU, 48GB RAM. The python version is 3.11, jax/jaxlib version is 0.4.7+cuda11+cudnn8.2, tensorflow version is 2.13.0.
It would be great if you can help me out there! Thx!