114 lines
4.3 KiB
Python
114 lines
4.3 KiB
Python
# -*- coding: utf-8 -*-
|
|
'''
|
|
@File : multi_pipeline.py
|
|
@Time : 2025/12/06 15:32:44
|
|
@Author : wty-yy
|
|
@Version : 1.0
|
|
@Blog : https://wty-yy.github.io/
|
|
@Desc : Multiprocessing Pipeline for Robogauge
|
|
'''
|
|
import yaml
|
|
import functools
|
|
import numpy as np
|
|
from tqdm import tqdm
|
|
import multiprocessing
|
|
from pathlib import Path
|
|
from copy import deepcopy
|
|
from itertools import product
|
|
from collections import defaultdict
|
|
|
|
from robogauge.tasks.pipeline.base_pipeline import BasePipeline
|
|
|
|
from robogauge.utils.task_register import task_register
|
|
from robogauge.utils.logger import logger
|
|
|
|
def run_single_process(args, data):
|
|
seed, base_mass, friction = data
|
|
local_args = deepcopy(args)
|
|
local_args.seed = seed
|
|
run_name = f"{local_args.run_name}_{seed}_baseMass{base_mass}_friction{friction}"
|
|
logger.create(
|
|
experiment_name=local_args.experiment_name,
|
|
run_name=run_name,
|
|
console_output=False
|
|
)
|
|
pipeline = task_register.make_pipeline(args=local_args, create_logger=False)
|
|
log_dir = pipeline.run()
|
|
return log_dir
|
|
|
|
class MultiPipeline:
|
|
def __init__(self, args):
|
|
self.args = args
|
|
self.seeds = args.seeds
|
|
self.frictions = args.frictions
|
|
self.base_masses = args.base_masses
|
|
self.num_processes = args.num_processes
|
|
logger.create(args.experiment_name+'_multi', args.run_name+'_multi')
|
|
|
|
def run(self):
|
|
logger.info(f"🚀 Starting Multi-Process Evaluation with {self.num_processes} processes.")
|
|
logger.info(f"🔢 Seeds: {self.seeds}, Frictions: {self.frictions}, Base masses: {self.base_masses}")
|
|
|
|
process_args = list(product(self.seeds, self.base_masses, self.frictions))
|
|
ctx = multiprocessing.get_context('spawn')
|
|
worker_func = functools.partial(run_single_process, self.args)
|
|
result_log_dirs = []
|
|
with ctx.Pool(processes=self.num_processes) as pool:
|
|
iterator = pool.imap_unordered(worker_func, process_args)
|
|
for log_dir in tqdm(iterator, total=len(process_args), desc="Evaluation"):
|
|
result_log_dirs.append(log_dir)
|
|
|
|
logger.info("✅ Multi-Process Evaluation Completed.")
|
|
self.aggregate_results(result_log_dirs)
|
|
|
|
def aggregate_results(self, log_dirs):
|
|
""" Process results.yaml from each log_dir """
|
|
all_results = []
|
|
all_yaml_paths = []
|
|
for path in log_dirs:
|
|
yaml_path = Path(path) / "results.yaml"
|
|
if not yaml_path.exists():
|
|
logger.warning(f"Results file not found: {yaml_path}, skipping.")
|
|
continue
|
|
with open(yaml_path, 'r') as file:
|
|
data = yaml.safe_load(file)
|
|
if data:
|
|
all_results.append(data)
|
|
all_yaml_paths.append(yaml_path)
|
|
if not all_results:
|
|
logger.error("No results to aggregate.")
|
|
return
|
|
yaml_paths_str = '\n'.join([str(p) for p in all_yaml_paths])
|
|
logger.info(
|
|
f"""\n{'='*20} Results Files {'='*20}\n"""
|
|
f"""{yaml_paths_str}\n"""
|
|
f"""{'='*56}"""
|
|
)
|
|
|
|
value_collections = defaultdict(lambda: defaultdict(list))
|
|
for result in all_results:
|
|
for goal, metrics in result.items():
|
|
if goal != 'summary':
|
|
continue
|
|
for metric, means in metrics.items():
|
|
for mean_name, mean_value in means.items():
|
|
value_collections[metric][mean_name].append(float(mean_value.split(' ')[0]))
|
|
|
|
summary = {}
|
|
for metric, means in value_collections.items():
|
|
summary[metric] = {}
|
|
for mean_name, values in means.items():
|
|
summary[metric][mean_name] = f"{float(np.mean(values)):.4f} ± {float(np.std(values)):.4f}"
|
|
|
|
save_path = logger.log_dir / "aggregated_results.yaml"
|
|
with open(save_path, 'w') as file:
|
|
yaml.dump(summary, file, allow_unicode=True)
|
|
logger.info("✅ Aggregated execution finished.")
|
|
logger.info(f"📁 Aggregated results saved to: {save_path}")
|
|
|
|
logger.info(
|
|
f"""\n{'='*20} Multi-Run Summary {'='*20}\n"""
|
|
f"""{yaml.dump(summary, allow_unicode=True)}"""
|
|
f"""{'='*60}"""
|
|
)
|