64 lines
2.3 KiB
Python
64 lines
2.3 KiB
Python
# -*- coding: utf-8 -*-
|
|
'''
|
|
@File : base_robot.py
|
|
@Time : 2025/11/27 15:53:57
|
|
@Author : wty-yy
|
|
@Version : 1.0
|
|
@Blog : https://wty-yy.github.io/
|
|
@Desc : Base Robot Class
|
|
'''
|
|
import torch
|
|
import numpy as np
|
|
|
|
from robogauge.utils.helpers import parse_path
|
|
from robogauge.utils.logger import logger
|
|
from robogauge.tasks.robots.base_robot_config import RobotConfig
|
|
from robogauge.tasks.simulator.sim_data import SimData
|
|
from robogauge.tasks.gauge.goal_data import GoalData
|
|
|
|
class BaseRobot:
|
|
def __init__(self, cfg: RobotConfig):
|
|
self.cfg = cfg
|
|
self.device = self.cfg.control.device
|
|
self.num_obs = cfg.control.num_observations
|
|
self.num_action = cfg.control.num_actions
|
|
self.control_type = cfg.control.control_type
|
|
self.p_gains = np.array(cfg.control.p_gains)
|
|
self.d_gains = np.array(cfg.control.d_gains)
|
|
model_path = parse_path(cfg.control.model_path)
|
|
logger.info(f"Loading robot model from '{model_path}'")
|
|
self.model = torch.jit.load(model_path).to(self.device)
|
|
self.model.eval()
|
|
|
|
def build_observation(self, sim_data: SimData, goal_data: GoalData) -> np.ndarray:
|
|
obs = np.zeros(self.num_obs, dtype=np.float32)
|
|
return obs
|
|
|
|
def get_action(self, obs: np.ndarray):
|
|
"""
|
|
Returns:
|
|
action: (num_action,) target joint positions/velocities/torques
|
|
p_gains: (num_action,) proportional gains for Mujoco PD controller
|
|
d_gains: (num_action,) derivative gains for Mujoco PD controller
|
|
control_type: 'P', 'V', or 'T' for position/velocity/torque control
|
|
"""
|
|
action = np.zeros(self.num_action, dtype=np.float32)
|
|
return action, self.p_gains, self.d_gains, self.control_type
|
|
|
|
def get_projected_gravity(quat):
|
|
""" Compute world frame gravity (0, 0, -1) projected into robot base frame.
|
|
Args:
|
|
quat: (4,) quaternion (w, x, y, z) from robot base to world frame
|
|
Returns:
|
|
projected_gravity: (3,) projected gravity vector in robot base frame
|
|
"""
|
|
qw, qx, qy, qz = quat
|
|
|
|
gravity_orientation = np.zeros(3)
|
|
|
|
gravity_orientation[0] = 2 * (-qz * qx + qw * qy)
|
|
gravity_orientation[1] = -2 * (qz * qy + qw * qx)
|
|
gravity_orientation[2] = 1 - 2 * (qw * qw + qz * qz)
|
|
|
|
return gravity_orientation
|