get_avg_reward.py 1.0 KB

123456789101112131415161718192021222324252627282930313233
  1. import argparse
  2. import json
  3. import browsergym.miniwob # noqa F401 register miniwob tasks as gym environments
  4. import gymnasium as gym
  5. parser = argparse.ArgumentParser(description='Calculate average reward.')
  6. parser.add_argument('output_path', type=str, help='path to output.jsonl')
  7. args = parser.parse_args()
  8. if __name__ == '__main__':
  9. env_ids = [
  10. id for id in gym.envs.registry.keys() if id.startswith('browsergym/miniwob')
  11. ]
  12. total_num = len(env_ids)
  13. print('Total number of tasks: ', total_num)
  14. total_reward = 0
  15. total_cost = 0
  16. actual_num = 0
  17. with open(args.output_path, 'r') as f:
  18. for line in f:
  19. data = json.loads(line)
  20. actual_num += 1
  21. total_cost += data['metrics']['accumulated_cost']
  22. total_reward += data['test_result']['reward']
  23. avg_reward = total_reward / total_num
  24. print('Avg Reward: ', avg_reward)
  25. avg_cost = total_cost / actual_num
  26. print('Avg Cost: ', avg_cost)
  27. print('Actual number of tasks finished: ', actual_num)