Action repeat for PyTorch training

[conglu1997]

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!

[DavidSlayback]

Had to dig into this recently myself! You should be able to set it with:

env = gym.make(gym_name, batch_size=num_envs, episode_length=episode_length, action_repeat=4)

Under the hood, that will propagate to the brax EpisodeWrapper