manual_control_matplotlib.py 849 B

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152
  1. import time
  2. import matplotlib.pyplot as plt
  3. import numpy as np
  4. import gym
  5. import gym_minigrid
  6. def key_handler(event):
  7. print('pressed', event.key)
  8. if event.key == 'escape':
  9. plt.close()
  10. return
  11. if event.key == 'left':
  12. env.step(env.actions.left)
  13. img = env.render('rgb_array')
  14. #img = np.zeros(shape=(256,256,3), dtype=np.uint8)
  15. imshow_obj.set_data(img)
  16. fig.canvas.draw()
  17. #plt.show()
  18. return
  19. env = gym.make('MiniGrid-Empty-8x8-v0')
  20. #env.step(env.actions.left)
  21. t0 = time.time()
  22. for i in range(1000):
  23. img = env.render('rgb_array')
  24. t1 = time.time()
  25. dt = int(1000 * (t1-t0))
  26. print(dt)
  27. print(img.shape)
  28. fig, ax = plt.subplots()
  29. fig.canvas.mpl_connect('key_press_event', key_handler)
  30. #plt.figure(num='gym-minigrid')
  31. imshow_obj = ax.imshow(img)
  32. plt.show()