There's another issue. The renderings are not in color. Here are some example gifs and the code used to produce them.

[jamesheald]

          There's another issue. The renderings are not in color. Here are some example gifs and the code used to produce them.

import jax.numpy as jnp
from brax import envs
from brax.io import image
from jax import random
import jax
from IPython.display import Image

env_name = "reacher"
backend = "generalized"
episode_length = 5

env = envs.get_environment(backend=backend, env_name=env_name)

state = jax.jit(env.reset)(random.PRNGKey(0))

cum_reward = 0
states = []
rollout = []
for step in range(1, episode_length + 1):

    rollout.append(state.pipeline_state)

    action = jnp.array([0]*env.action_size)
    state = jax.jit(env.step)(state, action)
    states.append(state.pipeline_state)
    cum_reward += state.reward

    if state.done:
        break

gif = Image(image.render(env.sys, rollout, fmt = 'gif'))
open('reacher.gif', 'wb').write(gif.data)

Originally posted by @jamesheald in https://github.com/google/brax/issues/460#issuecomment-1968938958

[btaba]

Hi @jamesheald , we'll add rgba to the envs, thanks for pointing this out!