How to improve tensorflow 2.0 code for policy gradient?
If possible, try to use just tensorflow functions (put the .numpy() part out, for example), and put a @tf.function decorator on top of train() function.

The role of @tf.function is that it transforms a whole function into a tensorflow op. The whole function will be execute an order of magnitude faster than a normal Python function.