robfiras / loco-mujoco

Imitation learning benchmark focusing on complex locomotion tasks using MuJoCo.
MIT License
475 stars 38 forks source link

Rendering Mode #14

Closed mertalbaba closed 4 months ago

mertalbaba commented 4 months ago

It seems like Humanoid environments does not allow rgb_array rendering mode. Is it possible to render the environment frames in a headless server?

robfiras commented 4 months ago

I am not sure if I understand what you would like to do. Training with visual observations is not supported in LocoMuJoCo. If you just want to record a video of a trained agent in headless mode, you can do so by passing a flag to the environment:

env = LocoEnv.make("Atlas", headless=True)

This will allow you to render a video on a headless server.

mertalbaba commented 4 months ago

What I was trying to do is recording videos of rollouts, and a common way to that is rendering gym environment with mode 'rgb_array', that returns rgb image when env.render is called. However 'rgb_array' rendering mode is not allowed in Loco-Mujoco tasks.

Can you please clarify what is the effect of setting headless=True? When I call render function of the environment with this initialization, does it return an rgb image? Or should I record the video through some wrapper?

robfiras commented 4 months ago

I think I understand your issue now. The headless tag is used to suppress opening an open-gl window while rendering in our viewer, which would raise an error on headless machines.

So then you are using the Gymansium interface for your training loop? We indeed did not implement rendering mode for the Gymanisum interface, but I can do that very quickly.

If you are using MushroomRL by yourself or just want to record some of the imitation learning algorithms in our examples, I can give further information on how to do the rendering in headless mode there.

mertalbaba commented 4 months ago

Yes, I am using Gymnasium interface. It would be really good if you can implement rendering mode as you said. Thanks a lot!

robfiras commented 4 months ago

I will implement it today.

robfiras commented 4 months ago

I added two render modes for the Gymansium interface, "human" and "rgb_array". rgb_array now returns an image, and supports headless rendering. Here is a minimal example:

import numpy as np
from loco_mujoco import LocoEnv
import gymnasium as gym

# create the environment and task
env = gym.make("LocoMujoco", env_name="UnitreeH1.run.real", render_mode="rgb_array")

action_dim = env.action_space.shape[0]

env.reset()
img = env.render()      # get the image
terminated = False
i = 0

while True:
    if i == 1000 or terminated:
        env.reset()
        i = 0
    action = np.random.randn(action_dim)
    nstate, reward, terminated, truncated, info = env.step(action)

    img = env.render()    # get the image 
    i += 1

Let me know if there is something missing for you and thanks for making me aware of this!

mertalbaba commented 4 months ago

Thanks a lot for a very fast solution!