Speed up the policy optimization by @tf.function

[ikamensh]

Using @tf.function allowed for at 2-3 times speedup of policy optimization (longest step) on 'Pendulum-v0' environment. This works by pre-compiling the tensorflow graph on first execution of the function. Please let me know if you need more careful speedup analysis.

[nrontsis]

@kyr-pol what do you think?

[kyr-pol]

That sounds good to me, could you test whether the change is compatible with the safety extension too (like the safe_cars_run script in the examples)? I think I had tried the tf.function decorator when moving the project to tensorflow 2.x, and I faced some issues, which might have been on my part or patched by now. If there are no issues there and the tests are successful, we can go ahead.

[ikamensh]

I don't have Mujoco license, so I can only test the safe_cars_run.py. From high-level examination it seems not to show much of a difference for two versions: see output per links

https://justpaste.it/5ozx8 vs https://justpaste.it/2ebz3

I must say I am not very sure how to read this output, but the final reward is same order of magnitude. Speedup is not that obvious in that environment.

[kyr-pol]

Great, that looks totally fine to me, I was mostly worried for tensorflow errors/crashes, this looks reasonable!

[gitAugust]

Using @tf.function will transfor EagerTensor to Tensor with can't iterate automatically, so there will be errors with safe_swimmer_run.py. Do you gays have some solutions?