Adding type annotations to train fn in train_ppo_lib.
Adding type annotations to train fn in train_ppo_lib.

PiperOrigin-RevId: 440443054
Change-Id: If7201d07bb1b3c5ea7720841b84277ef661e42cb

google-research/circuit_training