Parcourir la source

Added code to benchmark agent view rendering

Maxime Chevalier-Boisvert il y a 5 ans
Parent
commit
f336f97680
1 fichiers modifiés avec 17 ajouts et 4 suppressions
  1. 17 4
      benchmark.py

+ 17 - 4
benchmark.py

@@ -4,6 +4,7 @@ import time
 import argparse
 import gym_minigrid
 import gym
+from gym_minigrid.wrappers import *
 
 parser = argparse.ArgumentParser()
 parser.add_argument(
@@ -32,9 +33,21 @@ for i in range(args.num_frames):
     env.render('rgb_array')
 t1 = time.time()
 dt = t1 - t0
-dt_per_frame = (1000 * dt) / args.num_frames
 frames_per_sec = args.num_frames / dt
 
-print('reset time: {:.1f} ms'.format(reset_time))
-print('dt per frame: {:.1f} ms'.format(dt_per_frame))
-print('frames per sec: {:.0f}'.format(frames_per_sec))
+# Create an environment with an RGB agent observation
+env = gym.make(args.env_name)
+env = RGBImgPartialObsWrapper(env)
+env = ImgObsWrapper(env)
+
+# Benchmark rendering
+t0 = time.time()
+for i in range(args.num_frames):
+    obs, reward, done, info = env.step(0)
+t1 = time.time()
+dt = t1 - t0
+agent_view_fps = args.num_frames / dt
+
+print('Env reset time: {:.1f} ms'.format(reset_time))
+print('Rendering FPS : {:.0f}'.format(frames_per_sec))
+print('Agent view FPS: {:.0f}'.format(agent_view_fps))