191 lines
7.9 KiB
Python
191 lines
7.9 KiB
Python
# -*- coding: utf-8 -*-
|
|
'''
|
|
@File : base_gauge.py
|
|
@Time : 2025/11/27 15:55:19
|
|
@Author : wty-yy
|
|
@Version : 1.0
|
|
@Blog : https://wty-yy.github.io/
|
|
@Desc : Base Gauge for Robogauge
|
|
'''
|
|
import yaml
|
|
from typing import List
|
|
from pathlib import Path
|
|
from functools import partial
|
|
from collections import defaultdict
|
|
|
|
import numpy as np
|
|
from robogauge.utils.logger import logger
|
|
from robogauge.utils.helpers import class_to_dict, snake_to_pascal
|
|
|
|
from robogauge.tasks.robots import RobotConfig
|
|
from robogauge.tasks.gauge.base_gauge_config import BaseGaugeConfig
|
|
from robogauge.tasks.gauge.goal_data import GoalData, VelocityGoal, PositionGoal
|
|
from robogauge.tasks.simulator.sim_data import SimData
|
|
from robogauge.tasks.gauge.gauge_configs.terrain_levels_config import SEARCH_LEVELS_TERRAINS
|
|
|
|
from robogauge.tasks.gauge.goals import *
|
|
from robogauge.tasks.gauge.metrics import *
|
|
|
|
class BaseGauge:
|
|
def __init__(self,
|
|
cfg: BaseGaugeConfig,
|
|
robot_cfg: RobotConfig,
|
|
):
|
|
self.cfg = cfg
|
|
self.robot_cfg = robot_cfg
|
|
self.goals_cfg = class_to_dict(self.cfg.goals)
|
|
self.metrics_cfg = class_to_dict(self.cfg.metrics)
|
|
|
|
self.goal_str = "Init"
|
|
self.goal_idx = 0
|
|
self.goals: List[BaseGoal] = []
|
|
self.metrics: List[BaseMetric] = []
|
|
self.info = {'goal': [], 'metric': []}
|
|
self.results = {} # {'goal/sub_goal': {'metric': result}}
|
|
|
|
log_str = "Initialized Gauge with Goals 🎯 and Metrics 📊:\n"
|
|
for name, kwargs in self.goals_cfg.items():
|
|
if not kwargs['enabled']: continue
|
|
if name == 'max_velocity':
|
|
self.goals.append(MaxVelocityGoal(robot_cfg.control.control_dt, robot_cfg.commands, **kwargs))
|
|
log_str += f" - Max Velocity Goal: {kwargs}\n"
|
|
elif name == 'diagonal_velocity':
|
|
self.goals.append(DiagonalVelocityGoal(robot_cfg.control.control_dt, robot_cfg.commands, **kwargs))
|
|
log_str += f" - Diagonal Velocity Goal: {kwargs}\n"
|
|
elif name == 'target_pos_velocity':
|
|
self.goals.append(TargetPosVelocityGoal(robot_cfg.control.control_dt, cfg.backward, **kwargs))
|
|
log_str += f" - Target Position Velocity Goal: {kwargs}\n"
|
|
elif name == 'joystick':
|
|
self.goals.append(JoystickGoal(robot_cfg.commands, **kwargs))
|
|
log_str += f" - Joystick Goal: {kwargs}\n"
|
|
elif name == 'keyboard':
|
|
self.goals.append(KeyboardGoal(robot_cfg.commands, **kwargs))
|
|
log_str += f" - Keyboard Goal: {kwargs}\n"
|
|
else:
|
|
raise NotImplementedError(f"Goal '{name}' is not implemented in BaseGauge.")
|
|
self.info['goal'].append(name)
|
|
for name, enabled in self.metrics_cfg.items():
|
|
if not enabled: continue
|
|
if name in ['metric_dt']: continue
|
|
metric_class_name = f"{snake_to_pascal(name)}Metric"
|
|
metric_class = eval(metric_class_name)
|
|
self.metrics.append(metric_class(robot_cfg=robot_cfg, **self.metrics_cfg[name]))
|
|
log_str += f" - Metric: {name}\n"
|
|
self.info['metric'].append(name)
|
|
logger.info(log_str.strip())
|
|
|
|
if len(self.goals) == 0:
|
|
logger.warning("No goals have been configured for the Gauge. Exiting.")
|
|
else:
|
|
self.create_new_goal_logger()
|
|
|
|
def is_reset(self, sim_data: SimData) -> bool:
|
|
if self.goal_idx >= len(self.goals):
|
|
return False
|
|
return self.goals[self.goal_idx].is_reset(sim_data)
|
|
|
|
def is_done(self) -> bool:
|
|
if self.goal_idx >= len(self.goals):
|
|
self.save_results()
|
|
return True
|
|
return False
|
|
|
|
def create_new_goal_logger(self):
|
|
""" Create a new logger for new goal to metrics. """
|
|
if self.goal_idx >= len(self.goals) or not self.cfg.write_tensorboard:
|
|
return
|
|
logger.create_tensorboard(
|
|
self.robot_cfg.robot_name,
|
|
Path(self.robot_cfg.control.model_path).stem,
|
|
self.goals[self.goal_idx].name
|
|
)
|
|
|
|
def get_goal(self, sim_data: SimData) -> GoalData:
|
|
# goal = GoalData(
|
|
# goal_type='velocity',
|
|
# velocity_goal=VelocityGoal(
|
|
# ang_vel_yaw=-5.0,
|
|
# )
|
|
# )
|
|
goal_obj = self.goals[self.goal_idx]
|
|
if goal_obj.pre_get_goal(sim_data): # goal finished
|
|
self.switch_to_next_goal()
|
|
return None
|
|
goal = goal_obj.get_goal(sim_data)
|
|
|
|
now_goal_str = str(goal_obj)
|
|
if f"{goal_obj.count+1}_{now_goal_str}" != self.goal_str: # sub goal changed
|
|
self.goal_str = f"{goal_obj.count+1}_{now_goal_str}"
|
|
logger.info(f"New Goal [{self.goal_idx+1}/{len(self.goals)}] [{goal_obj.count+1}/{goal_obj.total}]: {self.goal_str}")
|
|
return goal
|
|
|
|
def reset_current_goal(self):
|
|
""" Reset the current goal """
|
|
if self.goal_idx >= len(self.goals):
|
|
return
|
|
goal_obj = self.goals[self.goal_idx]
|
|
goal_obj.reset_goal()
|
|
|
|
def switch_to_next_goal(self):
|
|
""" Switch to the next goal and log the metrics of the current goal. """
|
|
goal_obj = self.goals[self.goal_idx]
|
|
metrics = goal_obj.goal_mean_metrics
|
|
if hasattr(goal_obj, 'success'): # target position goal
|
|
metrics['success'] = {'mean': float(goal_obj.success)}
|
|
|
|
key = f"{goal_obj.name}"
|
|
self.results[key] = metrics
|
|
|
|
self.goal_idx += 1
|
|
self.create_new_goal_logger()
|
|
|
|
def update_metrics(self, sim_data: SimData, goal_data: GoalData):
|
|
if sim_data.n_step % int(self.cfg.metrics.metric_dt / sim_data.sim_dt + 1e-9) != 0:
|
|
return False
|
|
metrics_results = {}
|
|
for metric_name, metric_obj in zip(self.info['metric'], self.metrics):
|
|
val = metric_obj(sim_data, goal_data)
|
|
if metric_name not in ['visualization']:
|
|
metrics_results[metric_name] = val
|
|
self.goals[self.goal_idx].update_metrics(metrics_results)
|
|
return True
|
|
|
|
def reset_metrics(self):
|
|
for metric in self.metrics:
|
|
metric.reset()
|
|
|
|
def save_results(self):
|
|
""" Save the results to a yaml file. """
|
|
metrics = defaultdict(lambda: defaultdict(list))
|
|
for goal in self.results:
|
|
for metric_name, quantiles in self.results[goal].items():
|
|
for quantile, val in quantiles.items():
|
|
metrics[metric_name][quantile].append(val)
|
|
self.results['summary'] = {'quality_score': {}, 'terrain_quality_score': {}}
|
|
for metric_name, quantiles in metrics.items():
|
|
if metric_name not in self.results['summary']:
|
|
self.results['summary'][metric_name] = {}
|
|
for quantile, vals in quantiles.items():
|
|
mean = float(np.mean(vals))
|
|
std = float(np.std(vals))
|
|
self.results['summary'][metric_name][quantile] = f"{mean:.4f} ± {std:.4f}"
|
|
if metric_name == 'quality_score':
|
|
tqs = mean
|
|
if self.cfg.assets.terrain_name in SEARCH_LEVELS_TERRAINS:
|
|
tqs = 0.09 * (self.cfg.assets.terrain_level - 1) + 0.19 * mean
|
|
self.results['summary']['terrain_quality_score'][quantile] = f"{tqs:.4f} ± {std:.4f}"
|
|
|
|
save_path = Path(logger.log_dir) / "results.yaml"
|
|
self.results["terrain_name"] = self.cfg.assets.terrain_name
|
|
self.results["terrain_level"] = self.cfg.assets.terrain_level
|
|
|
|
with open(save_path, 'w', encoding='utf-8') as file:
|
|
yaml_str = yaml.dump(self.results, allow_unicode=True, sort_keys=False)
|
|
file.write(yaml_str)
|
|
logger.info(
|
|
f"""\n{'='*20} Goals and Metrics results {'='*20}\n"""
|
|
f"""{yaml_str}"""
|
|
f"""{'='*68}"""
|
|
)
|
|
logger.info(f"Saved metric results to {save_path}")
|