navsim/planning/script/grid_search_ensemble_unlog.py

import logging
import os
import pickle
from pathlib import Path

import pandas as pd
import torch
import copy

logger = logging.getLogger(__name__)

"""
pkl -> search params and calculation process
"""

import argparse

parser=argparse.ArgumentParser()
parser.add_argument('--pkl_path_vov', required=True)
parser.add_argument('--pkl_path_moe', required=True)
parser.add_argument('--pkl_path_davit', required=True)
parser.add_argument('--pkl_path_intern', required=True)

def main() -> None:
    args = parser.parse_args()
    pkl_path_vov = args.pkl_path_vov
    pkl_path_moe = args.pkl_path_moe
    pkl_path_davit = args.pkl_path_davit
    pkl_path_intern = args.pkl_path_intern

    # pkl_path_vov = f'{os.getenv("NAVSIM_EXP_ROOT")}/ensemble_navtest/vov_trainval_512x2048_epoch17.pkl'
    # pkl_path_moe = f'{os.getenv("NAVSIM_EXP_ROOT")}/ensemble_navtest/da+eva+vov_trainval_512x2048_epoch12.pkl'
    # pkl_path_davit = f'{os.getenv("NAVSIM_EXP_ROOT")}/ensemble_navtest/davit_trainval_256x1024_epoch16.pkl'

    predictions = {
        'vov': pickle.load(open(pkl_path_vov, 'rb')),
        'moe': pickle.load(open(pkl_path_moe, 'rb')),
        'davit': pickle.load(open(pkl_path_davit, 'rb')),
        'intern': pickle.load(open(pkl_path_intern, 'rb'))
    }
    models = ['vov', 'davit', 'moe']
    # prop_vov 0.5
    # prop_moe 0.1
    # prop_davit 0.4
    weights = {
        'vov': {
            'imi': 0.02,
            'noc': 0.7,
            'da': 0.1,
            'ttc': 5.0,
            'progress': 5.0,
            'comfort': 2.0,
            'tpc': 8.0
        },
        'moe': {
            'imi': 0.03,
            'noc': 0.001,
            'da': 0.024,
            'ttc': 5.0,
            'progress': 5.0,
            'comfort': 2.0,
            'tpc': 7.0
        },
        'davit': {
            'imi': 0.02,
            'noc': 0.6,
            'da': 0.5,
            'ttc': 5.0,
            'progress': 5.0,
            'comfort': 2.0,
            'tpc': 3.0
        },
        'intern': {
            'imi': 0.01,
            'noc': 0.1,
            'da': 0.4,
            'ttc': 5.0,
            'progress': 5.0,
            'comfort': 2.0,
            'tpc': 1.0
        }
    }

    metric_keys = {
        'imi': [],
        'noc': [],
        'da': [],
        'ttc': [],
        'progress': [],
        'comfort': []
    }
    tensor_predictions = {
        model_name: copy.deepcopy(metric_keys) for model_name in models
    }

    navtest_scores = pickle.load(
        open(f'{os.getenv("NAVSIM_TRAJPDM_ROOT")}/vocab_score_full_8192_navtest/navtest.pkl', 'rb')
    )

    pdm_scores = []
    total_scene_cnt = len(navtest_scores)
    print(f'total_scene_cnt: {total_scene_cnt}')
    for token, v in navtest_scores.items():
        pdm_scores.append(torch.from_numpy(v['total'][None]).cuda())
        for metric_k in metric_keys:
            for model in models:
                tensor_predictions[model][metric_k].append(torch.from_numpy(predictions[model][token][metric_k][None]).cuda())

    pdm_scores = torch.cat(pdm_scores, 0).contiguous()
    for metric_k in metric_keys:
        for model in models:
            tensor_predictions[model][metric_k] = torch.cat(tensor_predictions[model][metric_k], 
                                                            0).contiguous()

    proportions_vov = [0.5]
    proportions_moe = [0.1]
    proportions_davit = [0.4]
    proportions_intern = [1]
    highest_info = {
        'score': -100,
    }
    for prop_vov in proportions_vov:
        for prop_moe in proportions_moe:
            for prop_davit in proportions_davit:
                for prop_intern in proportions_intern:
                    scores = 0.0
                    for model in models:
                        tmp_score = (weights[model]['imi'] * tensor_predictions[model]['imi'] + 
                                    weights[model]['noc'] * tensor_predictions[model]['noc'] + 
                                    weights[model]['da'] * tensor_predictions[model]['da'] + 
                                    weights[model]['tpc'] * (
                                        weights[model]['ttc'] * tensor_predictions[model]['ttc'].exp() +
                                        weights[model]['progress'] * tensor_predictions[model]['progress'].exp() +
                                        weights[model]['comfort'] * tensor_predictions[model]['comfort'].exp()
                                    ).log()
                        )
                        if model == 'vov':
                            scores += tmp_score * prop_vov
                        elif model == 'moe':
                            scores += tmp_score * prop_moe
                        elif model == 'davit':
                            scores += tmp_score * prop_davit
                        elif model == 'intern':
                            scores += tmp_score * prop_intern
                        else:
                            raise ValueError('what model?')

                    pdm_score = pdm_scores[
                        torch.arange(total_scene_cnt, device=pdm_scores.device),
                        scores.argmax(-1)
                    ]

                    pdm_score = pdm_score.mean().item()
                    print(f'vov: {prop_vov}, moe: {prop_moe}, davit: {prop_davit}, intern: {prop_intern} score: {pdm_score}')
                    if pdm_score > highest_info['score']:
                        highest_info['score'] = pdm_score
                        highest_info['prop_vov'] = prop_vov
                        highest_info['prop_moe'] = prop_moe
                        highest_info['prop_davit'] = prop_davit
                        highest_info['prop_intern'] = prop_intern
                
   
    for k, v in highest_info.items():
        print(k, v)


if __name__ == "__main__":
    with torch.no_grad():
        main()