Kenji

We currently use TF agents in our Acme example. JAX has received more attention lately so should probably encourage researchers to use that instead.

This is a simple change to simplify engine.py to avoid duplicating some 5 lines in a if-else condition.