blob: d05a3abf647add0b7306d97f866cb934762d6b21 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
|
import gym
import time
import matplotlib.pyplot as plt
from matplotlib import animation
env_name = 'CartPole-v0'
env = gym.make(env_name)
state = env.reset()
done = False
total_score = 0
frames = []
while not done:
frames.append(env.render(mode='rgb_array'))
action = env.action_space.sample()
observation, reward, done, info = env.step(action)
total_score += reward
time.sleep(0.1)
def display_frame_as_gif(frames):
plt.figure(figsize=(frames[0].shape[1]/72, frames[0].shape[0]/72), dpi=72)
patch = plt.imshow(frames[0])
plt.axis('off')
def animate(i):
patch.set_data(frames[i])
anim = animation.FuncAnimation(plt.gcf(), animate, frames=len(frames), interval=50)
anim.save('movie_cartpole.gif')
env.close()
print(len(frames), frames[0].shape)
# display_frame_as_gif(frames)
|
