Last active
January 7, 2021 15:54
-
-
Save DanielTakeshi/fec9a5cd957eb05b04b6d06a16cc88ae to your computer and use it in GitHub Desktop.
How to make a GIF from MuJoCo environment observations programmatically.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import gym | |
| import numpy as np | |
| import imageio | |
| # Make the environment and the initial observation `o`. | |
| env_name = 'Walker2d-v3' | |
| env = gym.make(env_name) | |
| obs_dim = env.observation_space.shape | |
| act_dim = env.action_space.shape[0] | |
| act_limit = env.action_space.high[0] | |
| o = env.reset() | |
| # Book-keeping. | |
| num_episodes = 3 | |
| ep_ret = 0 | |
| ep_len = 0 | |
| ep_done = 0 | |
| ep_obs = [] | |
| # For video / GIF. | |
| dur = 0.01 | |
| width = 250 | |
| height = 200 | |
| while ep_done < num_episodes: | |
| obs = env.render(mode='rgb_array', width=width, height=height) | |
| assert obs.shape == (height, width, 3), obs.shape # height first! | |
| ep_obs.append(obs) | |
| # Take action, step into environment, etc. | |
| a = np.random.randn(act_dim) | |
| a = np.clip(a, -act_limit, act_limit) | |
| o, r, d, _ = env.step(a) | |
| ep_ret += r | |
| ep_len += 1 | |
| if d: | |
| # Form GIF. imageio should read from numpy: https://imageio.github.io/ | |
| print(f'Episode {ep_done}, cum. return: {ep_ret:0.1f}, length: {ep_len}.') | |
| ep_name = f'ep_{env_name}_{str(ep_done).zfill(2)}_dur_{dur}_len_{str(ep_len).zfill(3)}.gif' | |
| with imageio.get_writer(ep_name, mode='I', duration=dur) as writer: | |
| for obs_np in ep_obs: | |
| writer.append_data(obs_np) | |
| # Reset information. | |
| o = env.reset() | |
| ep_ret = 0 | |
| ep_len = 0 | |
| ep_done += 1 | |
| ep_obs = [] |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Mostly installed from spinningup conda env.
Note: I do NOT do
export LD_PRELOAD=...as in some solutions for getting MuJoCo to work. That will cause problems with RGB array mode.