71 lines
2.4 KiB
Python
71 lines
2.4 KiB
Python
import numpy as np
|
|
|
|
from robogauge.tasks.robots import RobotConfig
|
|
from robogauge.tasks.simulator.sim_data import SimData
|
|
|
|
from robogauge.utils.logger import logger
|
|
|
|
def example_metric(
|
|
sim_data: SimData,
|
|
robot_cfg: RobotConfig,
|
|
**kwargs
|
|
) -> float:
|
|
""" An example metric function. """
|
|
value = 0.0
|
|
# Compute some metric value based on sim_data and robot_cfg
|
|
logger.log(value, 'example_metric', step=sim_data.n_step)
|
|
return value
|
|
|
|
def dof_limits_metric(
|
|
sim_data: SimData,
|
|
robot_cfg: RobotConfig,
|
|
soft_dof_limit_ratio: float = 0.9,
|
|
dof_names: list = None,
|
|
**kwargs
|
|
) -> float:
|
|
""" Metric to log DOF limit violations. """
|
|
values = []
|
|
for i in range(len(sim_data.proprio.joint.limits)):
|
|
lower_limit = sim_data.proprio.joint.limits[i, 0]
|
|
upper_limit = sim_data.proprio.joint.limits[i, 1]
|
|
dof_range = upper_limit - lower_limit
|
|
soft_lower_limit = lower_limit + (1 - soft_dof_limit_ratio) * dof_range
|
|
soft_upper_limit = upper_limit - (1 - soft_dof_limit_ratio) * dof_range
|
|
|
|
pos = sim_data.proprio.joint.pos[i]
|
|
dof_name = sim_data.proprio.joint.names[i]
|
|
value = 0
|
|
if pos < soft_lower_limit:
|
|
value = soft_lower_limit - pos
|
|
elif pos > soft_upper_limit:
|
|
value = pos - soft_upper_limit
|
|
value /= dof_range # Normalize by DOF range
|
|
logger.log(value, f'dof_limits/{dof_name}', step=sim_data.n_step)
|
|
if dof_names is not None:
|
|
for use_name in dof_names:
|
|
if use_name in dof_name:
|
|
values.append(value)
|
|
else:
|
|
values.append(value)
|
|
rms_value = 1 - np.sqrt(np.mean(np.square(values)))
|
|
logger.log(1 - rms_value, f'dof_limits/rms', step=sim_data.n_step)
|
|
return rms_value
|
|
|
|
def visualization_metric(
|
|
sim_data: SimData,
|
|
robot_cfg: RobotConfig,
|
|
dof_force: bool = False,
|
|
dof_pos: bool = False,
|
|
**kwargs,
|
|
):
|
|
""" Metric to visualize various robot states in the simulator. """
|
|
for i in range(len(sim_data.proprio.joint.force)):
|
|
name = sim_data.proprio.joint.names[i]
|
|
if dof_force:
|
|
force = sim_data.proprio.joint.force[i]
|
|
logger.log(force, f'dof_force/{name}', step=sim_data.n_step)
|
|
if dof_pos:
|
|
pos = sim_data.proprio.joint.pos[i]
|
|
logger.log(pos, f'dof_pos/{name}', step=sim_data.n_step)
|
|
return 0.0
|