search

google / brax

Massively parallel rigidbody physics simulation on accelerator hardware.
Apache License 2.0
2.38k stars 257 forks source link

Action repeat for PyTorch training #162

Open conglu1997 opened 2 years ago

conglu1997 commented 2 years ago

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 commented 2 years ago

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