Integrate with Gym's RecordVideo wrapper.

[vwxyzjn]

This PR adds the meta-data tag that Gym's RecordVideo is looking for.

  metadata = {
    'render.modes': ['human', 'rgb_array'],
    'video.frames_per_second' : 24
  }

Additionally, in the case of a vector environment, the render function would only render the first sub environment for speed reasons.

You can run the following to give it a try:

import gym
import torch

from brax.envs.to_torch import JaxToTorchWrapper
from brax.envs import create_gym_env

dummy_value = torch.ones(1, device="cpu")
envs = create_gym_env("halfcheetah", batch_size=2)
envs.is_vector_env = True
envs = JaxToTorchWrapper(envs, device="cpu")
envs = gym.wrappers.RecordVideo(envs, f"videos/halfcheetah", video_length=50)

obs = envs.reset()
for _ in range(100):
    obs, reward, done, info = envs.step(envs.action_space.sample())
    if done[0]:
        break

[erikfrey]

Thanks for doing this! Two questions for you.

[erikfrey]

Should this be scaled according to the environment dt?

[vwxyzjn]

Sorry I don't quite follow, what dt?

[erikfrey]

When you step an environment, a certain amount of time passes, for example in Ant, each step is 0.05 seconds:

https://github.com/google/brax/blob/main/brax/envs/ant.py#L388

For the video speed to match the speed in the world, you would want video.frames_per_second to be 1 / 0.05 = 20

[vwxyzjn]

That makes sense. Thanks for the pointer. Just fixed it in the latest commit.

[erikfrey]

Looks great - thank you!