How to show episode in rendered openAI gym environment

user_9470543 0
Viewed 1962

If we look at the previews of the environments, they show the episodes increasing in the animation on the bottom right corner. https://gym.openai.com/envs/CartPole-v1/ .Is there a command to explicitly show that?

openai-gym
1 Answers

I don't think there is a command to do that directly available in OpenAI, but I've written some code that you can probably adapt to your purposes. This is the end result:

enter image description here

These is how I achieve the end result:

  1. For each step, you obtain the frame with env.render(mode='rgb_array')
  2. You convert the frame (which is a numpy array) into a PIL image
  3. You write the episode name on top of the PIL image using utilities from PIL.ImageDraw (see the function _label_with_episode_number in the code snippet).
  4. You save the labeled image into a list of frames.
  5. You render the list of frames as a GIF using matplotlib utilities.

Here is the code I wrote for obtaining a GIF of the behavior of a random agent with the Episode number displayed in the top left corner of each frame:

import os
import imageio
import numpy as np
from PIL import Image
import PIL.ImageDraw as ImageDraw
import matplotlib.pyplot as plt    


def _label_with_episode_number(frame, episode_num):
    im = Image.fromarray(frame)

    drawer = ImageDraw.Draw(im)

    if np.mean(im) < 128:
        text_color = (255,255,255)
    else:
        text_color = (0,0,0)
    drawer.text((im.size[0]/20,im.size[1]/18), f'Episode: {episode_num+1}', fill=text_color)

    return im


def save_random_agent_gif(env):
    frames = []
    for i in range(5):
        state = env.reset()        
        for t in range(500):
            action = env.action_space.sample()

            frame = env.render(mode='rgb_array')
            frames.append(_label_with_episode_number(frame, episode_num=i))

            state, _, done, _ = env.step(action)
            if done:
                break

    env.close()

    imageio.mimwrite(os.path.join('./videos/', 'random_agent.gif'), frames, fps=60)


env = gym.make('CartPole-v1')
save_random_agent_gif(env)

You can find a working version of the code here: https://github.com/RishabhMalviya/dqn_experiments/blob/master/train_and_visualize.py#L10

user_35958430
Related