How to use `env.render()` to visualize environment transitions frame by frame?

[dantp-ai]

rng = jax.random.PRNGKey(0)
rng, key_reset, key_act, key_step = jax.random.split(rng, 4)

# Instantiate the environment & its settings.
env, env_params = gymnax.make('FourRooms-misc')

# Reset the environment.
obs, state = env.reset(key_reset, env_params)

while True:
    # Sample a random action.
    action = env.action_space(env_params).sample(key_act)

    # Perform the step transition.
    n_obs, n_state, reward, done, _ = env.step(key_step, state, action, env_params)
    print(f"action: {action}, reward: {reward}, done: {done}")
    env.render(n_state, env_params)

I tried this simple example. The matplotlib figure is never shown. It is stuck and never opens.

Follow-up: Why is gymnax using matplotlib for rendering and not pygame (as gym) ?

[dantp-ai]

P.S. I understand that I can use the Visualizer to animate a sequence of transitions, but that seems like an overhead when you just want to get the agent's view during interaction with the environment.