75 lines
2.8 KiB
Python
75 lines
2.8 KiB
Python
|
import argparse
|
||
|
|
import pickle
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
from scenarionet_training.scripts.train_nuplan import config
|
||
|
|
from scenarionet_training.train_utils.callbacks import DrivingCallbacks
|
||
|
|
from scenarionet_training.train_utils.multi_worker_PPO import MultiWorkerPPO
|
||
|
|
from scenarionet_training.train_utils.utils import initialize_ray
|
||
|
|
|
||
|
|
|
||
|
|
class NumpyEncoder(json.JSONEncoder):
|
||
|
|
def default(self, obj):
|
||
|
|
if isinstance(obj, np.ndarray):
|
||
|
|
return obj.tolist()
|
||
|
|
elif isinstance(obj, np.int32):
|
||
|
|
return int(obj)
|
||
|
|
elif isinstance(obj, np.int64):
|
||
|
|
return int(obj)
|
||
|
|
return json.JSONEncoder.default(self, obj)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument("--start_index", type=int, default=0)
|
||
|
|
parser.add_argument("--ckpt_path", type=str, required=True)
|
||
|
|
parser.add_argument("--database_path", type=str, required=True)
|
||
|
|
parser.add_argument("--id", type=str, default="")
|
||
|
|
parser.add_argument("--num_scenarios", type=int, default=5000)
|
||
|
|
parser.add_argument("--num_workers", type=int, default=10)
|
||
|
|
parser.add_argument("--horizon", type=int, default=600)
|
||
|
|
parser.add_argument("--allowed_more_steps", type=int, default=50)
|
||
|
|
parser.add_argument("--max_lateral_dist", type=int, default=2.5)
|
||
|
|
parser.add_argument("--overwrite", action="store_true")
|
||
|
|
|
||
|
|
args = parser.parse_args()
|
||
|
|
file = "eval_{}_{}_{}".format(args.id, os.path.basename(args.ckpt_path), os.path.basename(args.database_path))
|
||
|
|
if os.path.exists(file) and not args.overwrite:
|
||
|
|
raise FileExistsError("Please remove {} or set --overwrite".format(file))
|
||
|
|
initialize_ray(test_mode=True, num_gpus=1)
|
||
|
|
|
||
|
|
config["callbacks"] = DrivingCallbacks
|
||
|
|
config["evaluation_num_workers"] = args.num_workers
|
||
|
|
config["evaluation_num_episodes"] = args.num_scenarios
|
||
|
|
config["metrics_smoothing_episodes"] = args.num_scenarios
|
||
|
|
config["custom_eval_function"] = None
|
||
|
|
config["num_workers"] = 0
|
||
|
|
config["evaluation_config"]["env_config"].update(dict(
|
||
|
|
start_scenario_index=args.start_index,
|
||
|
|
num_scenarios=args.num_scenarios,
|
||
|
|
sequential_seed=True,
|
||
|
|
store_map=False,
|
||
|
|
store_data=False,
|
||
|
|
allowed_more_steps=args.allowed_more_steps,
|
||
|
|
# no_map=True,
|
||
|
|
max_lateral_dist=args.max_lateral_dist,
|
||
|
|
curriculum_level=1, # disable curriculum
|
||
|
|
target_success_rate=1,
|
||
|
|
horizon=args.horizon,
|
||
|
|
episodes_to_evaluate_curriculum=args.num_scenarios,
|
||
|
|
data_directory=args.database_path,
|
||
|
|
use_render=False))
|
||
|
|
|
||
|
|
trainer = MultiWorkerPPO(config)
|
||
|
|
trainer.restore(args.ckpt_path)
|
||
|
|
|
||
|
|
ret = trainer._evaluate()["evaluation"]
|
||
|
|
with open(file + ".json", "w") as f:
|
||
|
|
json.dump(ret, f, cls=NumpyEncoder)
|
||
|
|
|
||
|
|
with open(file + ".pkl", "wb+") as f:
|
||
|
|
pickle.dump(ret, f)
|