39 lines
1.3 KiB
Python
39 lines
1.3 KiB
Python
# -*- coding: utf-8 -*-
|
|
'''
|
|
@File : visualization.py
|
|
@Time : 2025/12/18 20:19:09
|
|
@Author : wty-yy
|
|
@Version : 1.0
|
|
@Blog : https://wty-yy.github.io/
|
|
@Desc : Visualization Metric Implementation
|
|
'''
|
|
from robogauge.tasks.robots import RobotConfig
|
|
from robogauge.tasks.gauge.metrics.base_metric import BaseMetric, SimData, GoalData
|
|
|
|
from robogauge.utils.logger import logger
|
|
|
|
class VisualizationMetric(BaseMetric):
|
|
""" Metric to visualize various robot states in the simulator. """
|
|
name = 'visualization_metric'
|
|
|
|
def __init__(self,
|
|
robot_cfg: RobotConfig,
|
|
dof_torque: bool = False,
|
|
dof_pos: bool = False,
|
|
**kwargs
|
|
):
|
|
super().__init__(robot_cfg)
|
|
self.dof_torque = dof_torque
|
|
self.dof_pos = dof_pos
|
|
|
|
def __call__(self, sim_data: SimData, goal_data: GoalData) -> float:
|
|
for i in range(len(sim_data.proprio.joint.torque)):
|
|
name = sim_data.proprio.joint.names[i]
|
|
if self.dof_torque:
|
|
torque = sim_data.proprio.joint.torque[i]
|
|
logger.log(torque, f'dof_torque/{name}', step=sim_data.n_step)
|
|
if self.dof_pos:
|
|
pos = sim_data.proprio.joint.pos[i]
|
|
logger.log(pos, f'dof_pos/{name}', step=sim_data.n_step)
|
|
return 0.0
|