Action repeat for PyTorch training

The current PyTorch training routine in https://colab.research.google.com/github/google/brax/blob/main/notebooks/training_torch.ipynb lacks an option for `action_repeat` to be set. From some testing of the Torch version, it seems this is mandatory to learn things efficiently for `reacher` (in the default JAX version, `action_repeat=4`), Would it be possible for this to be added in?

Thanks so much!