-
-
Save JackBosco/2a58bedd217c69d0bce5b7d2273a3907 to your computer and use it in GitHub Desktop.
Save OpenAI Gym renders as GIFS
This file contains 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
from matplotlib import animation | |
import matplotlib.pyplot as plt | |
import gym | |
import collections | |
import os | |
""" | |
Ensure you have imagemagick installed with | |
sudo apt-get install imagemagick | |
Open file in CLI with: | |
xgd-open <filelname> | |
""" | |
#the list of frames is a named tuple with the matrix, action, and reward at each index | |
Frame = collections.namedtuple('Frame', ['image', 'action', 'reward']) | |
def save_frames_as_gif(frames, path='./', filename='gym_animation.gif'): | |
#Mess with this to change frame size | |
plt.figure(figsize=(frames[0].image.shape[1] / 72.0, frames[0].image.shape[0] / 72.0), dpi=72) | |
#allows for the extra display | |
ax = plt.gcf().add_subplot() | |
plt.axis('off') | |
patch = plt.imshow(frames[0].image) | |
#configure extra display | |
at = ax.text(x=0.05, y=0.5, s=f"Action: {frames[0].action}\nReward: {frames[0].reward}", fontsize=18 \ | |
, bbox={'facecolor': 'grey', 'alpha': 0.5, 'pad': 10}, fontweight='black', fontfamily='serif') | |
#added extra display to animate method | |
def animate(i): | |
patch.set_data(frames[i].image) | |
at.set_text(s=f"Action: {frames[i].action}\nReward: {frames[i].reward}") | |
anim = animation.FuncAnimation(plt.gcf(), animate, frames = len(frames), interval=100) | |
anim.save(path + filename, writer='imagemagick', fps=30) | |
def main(): | |
#Make gym env | |
env = gym.make('CartPole-v1') | |
#Run the env | |
observation = env.reset() | |
frames = [] | |
printAction = ['←', '→'] | |
for t in range(1000): | |
#Render to frames buffer | |
action = env.action_space.sample() | |
frames.append(Frame(env.render(mode='rgb_array'), printAction[int(action)], t)) | |
_, _, done, _ = env.step(action) | |
if done: | |
break | |
env.close() | |
#handles filename collisions | |
path='.' + os.sep | |
filename='gym_animation.gif' | |
n = 0 | |
while filename in os.listdir(path): | |
n += 1 | |
filename = 'gym_animation' + str(n) + '.gif' | |
save_frames_as_gif(frames, filename=filename, path=path) | |
if __name__ == '__main__': | |
main() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
I added a display that gives the current action and reward at each step