Go2Py_SIM/examples/execute_policy.py

51 lines
1.6 KiB
Python

from Go2Py.robot.interface.ros2 import GO2Real, ros2_init, ROS2ExecutorManager
import time
ros2_init()
import numpy as np
for i in range(10000):
q = np.zeros(12)
dq = np.zeros(12)
kp = np.zeros(12)
kd = np.zeros(12)
tau = np.zeros(12)
tau[0] = -0.8
robot.setCommands(q, dq, kp, kd, tau)
time.sleep(0.01)
def skew_symm(_v):
a, b, c = _v(0), _v(1), _v(2)
return np.array([[0, -c, b], [c, 0, -a], [-b, a, 0]])
def quat2rot(_q):
q0 = _q(3)
qv = _q[0:3]
return (2 * q0**2 - 1) * np.eye(3) + 2 * q0 + skew_symm(qv) + 2 * qv @ qv.T
class ExecutePolicy():
def __init__(self):
self.initClass()
def initClass(self):
self.robot = GO2Real(mode='lowlevel')
self.ros2_exec_manager = ROS2ExecutorManager()
self.ros2_exec_manager.add_node(self.robot)
self.ros2_exec_manager.start()
def getObservation(self):
self.js = self.robot.getJointStates()
self.imu = self.robot.getIMU()
# Calculate projected gravity vector as a torch array
quat = self.imu['quat']
R = quat2rot(self.imu['quat'])
# Calculate current dof_pos as torch array
self.obs_buf = torch.cat((self.projected_gravity,
(self.dof_pos[:, :self.num_actuated_dof] - self.default_dof_pos[:,
:self.num_actuated_dof]) * self.obs_scales.dof_pos,
self.dof_vel[:, :self.num_actuated_dof] * self.obs_scales.dof_vel,
self.actions
), dim=-1)