Rendering OpenAi Gym in Google Colaboratory

openai-gym-rendering-colab.py

# Install gym dependencies
!apt-get update > /dev/null 2>&1
!apt-get install python-opengl -y > /dev/null 2>&1
!apt install xvfb -y --fix-missing > /dev/null 2>&1
!apt-get install ffmpeg > /dev/null 2>&1
!apt-get install x11-utils > /dev/null 2>&1

# Install rendering environment
!pip install pyvirtualdisplay > /dev/null 2>&1
!pip install piglet > /dev/null 2>&1

# Activate virtual display
from pyvirtualdisplay import Display
display = Display(visible=0, size=(1400, 900))
display.start()

# This code creates a virtual display to draw game images on. 
# If running locally, just ignore it
import os
if type(os.environ.get('DISPLAY')) is not str or \
        len(os.environ.get('DISPLAY')) == 0:
    !bash ../xvfb start
    %env DISPLAY=:1
                
"""
Utility functions to enable video recording of gym environment and displaying it
To enable video, just do "env = wrap_env(env)""
"""


import io
import glob
import base64
from IPython.display import HTML
from IPython import display as ipythondisplay

def show_video():
    mp4list = glob.glob('video/*.mp4')
    if len(mp4list) > 0:
        mp4 = mp4list[0]
        video = io.open(mp4, 'r+b').read()
        encoded = base64.b64encode(video)

        content = ipythondisplay.display(HTML(data='''
        <video alt="test" autoplay loop controls style="height: 400px;">
            <source src="data:video/mp4;base64,{0}" type="video/mp4" />
        </video>
        '''.format(encoded.decode('ascii'))))
    else: 
        print("Couldn't find video")

def wrap_env(env):
    env = gym.wrappers.Monitor(env, './video', force=True)
    return env