From 8210615f15ceec2886c7d30015ae3891aa3b68d1 Mon Sep 17 00:00:00 2001 From: chzhang Date: Tue, 6 Dec 2022 23:31:38 +0800 Subject: rl env render save --- rl/tutorials/02_env_rendering.py | 40 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) create mode 100644 rl/tutorials/02_env_rendering.py (limited to 'rl/tutorials/02_env_rendering.py') diff --git a/rl/tutorials/02_env_rendering.py b/rl/tutorials/02_env_rendering.py new file mode 100644 index 0000000..d05a3ab --- /dev/null +++ b/rl/tutorials/02_env_rendering.py @@ -0,0 +1,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) + -- cgit v1.2.3