302 lines
13 KiB
Python
302 lines
13 KiB
Python
|
import time
|
||
|
|
|
||
|
|
import lcm
|
||
|
|
import numpy as np
|
||
|
|
import torch
|
||
|
|
import cv2
|
||
|
|
|
||
|
|
from go1_gym_deploy.lcm_types.pd_tau_targets_lcmt import pd_tau_targets_lcmt
|
||
|
|
|
||
|
|
lc = lcm.LCM("udpm://239.255.76.67:7667?ttl=255")
|
||
|
|
|
||
|
|
|
||
|
|
def class_to_dict(obj) -> dict:
|
||
|
|
if not hasattr(obj, "__dict__"):
|
||
|
|
return obj
|
||
|
|
result = {}
|
||
|
|
for key in dir(obj):
|
||
|
|
if key.startswith("_") or key == "terrain":
|
||
|
|
continue
|
||
|
|
element = []
|
||
|
|
val = getattr(obj, key)
|
||
|
|
if isinstance(val, list):
|
||
|
|
for item in val:
|
||
|
|
element.append(class_to_dict(item))
|
||
|
|
else:
|
||
|
|
element = class_to_dict(val)
|
||
|
|
result[key] = element
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
class LCMAgent():
|
||
|
|
def __init__(self, cfg, se, command_profile):
|
||
|
|
if not isinstance(cfg, dict):
|
||
|
|
cfg = class_to_dict(cfg)
|
||
|
|
self.cfg = cfg
|
||
|
|
self.se = se
|
||
|
|
self.command_profile = command_profile
|
||
|
|
|
||
|
|
self.dt = self.cfg["control"]["decimation"] * self.cfg["sim"]["dt"]
|
||
|
|
self.timestep = 0
|
||
|
|
|
||
|
|
self.num_obs = self.cfg["env"]["num_observations"]
|
||
|
|
self.num_envs = 1
|
||
|
|
self.num_privileged_obs = self.cfg["env"]["num_privileged_obs"]
|
||
|
|
self.num_actions = self.cfg["env"]["num_actions"]
|
||
|
|
self.num_commands = self.cfg["commands"]["num_commands"]
|
||
|
|
self.device = 'cpu'
|
||
|
|
|
||
|
|
if "obs_scales" in self.cfg.keys():
|
||
|
|
self.obs_scales = self.cfg["obs_scales"]
|
||
|
|
else:
|
||
|
|
self.obs_scales = self.cfg["normalization"]["obs_scales"]
|
||
|
|
|
||
|
|
self.commands_scale = np.array(
|
||
|
|
[self.obs_scales["lin_vel"], self.obs_scales["lin_vel"],
|
||
|
|
self.obs_scales["ang_vel"], self.obs_scales["body_height_cmd"], 1, 1, 1, 1, 1,
|
||
|
|
self.obs_scales["footswing_height_cmd"], self.obs_scales["body_pitch_cmd"],
|
||
|
|
# 0, self.obs_scales["body_pitch_cmd"],
|
||
|
|
self.obs_scales["body_roll_cmd"], self.obs_scales["stance_width_cmd"],
|
||
|
|
self.obs_scales["stance_length_cmd"], self.obs_scales["aux_reward_cmd"], 1, 1, 1, 1, 1, 1
|
||
|
|
])[:self.num_commands]
|
||
|
|
|
||
|
|
|
||
|
|
joint_names = [
|
||
|
|
"FL_hip_joint", "FL_thigh_joint", "FL_calf_joint",
|
||
|
|
"FR_hip_joint", "FR_thigh_joint", "FR_calf_joint",
|
||
|
|
"RL_hip_joint", "RL_thigh_joint", "RL_calf_joint",
|
||
|
|
"RR_hip_joint", "RR_thigh_joint", "RR_calf_joint", ]
|
||
|
|
self.default_dof_pos = np.array([self.cfg["init_state"]["default_joint_angles"][name] for name in joint_names])
|
||
|
|
try:
|
||
|
|
self.default_dof_pos_scale = np.array([self.cfg["init_state"]["default_hip_scales"], self.cfg["init_state"]["default_thigh_scales"], self.cfg["init_state"]["default_calf_scales"],
|
||
|
|
self.cfg["init_state"]["default_hip_scales"], self.cfg["init_state"]["default_thigh_scales"], self.cfg["init_state"]["default_calf_scales"],
|
||
|
|
self.cfg["init_state"]["default_hip_scales"], self.cfg["init_state"]["default_thigh_scales"], self.cfg["init_state"]["default_calf_scales"],
|
||
|
|
self.cfg["init_state"]["default_hip_scales"], self.cfg["init_state"]["default_thigh_scales"], self.cfg["init_state"]["default_calf_scales"]])
|
||
|
|
except KeyError:
|
||
|
|
self.default_dof_pos_scale = np.ones(12)
|
||
|
|
self.default_dof_pos = self.default_dof_pos * self.default_dof_pos_scale
|
||
|
|
|
||
|
|
self.p_gains = np.zeros(12)
|
||
|
|
self.d_gains = np.zeros(12)
|
||
|
|
for i in range(12):
|
||
|
|
joint_name = joint_names[i]
|
||
|
|
found = False
|
||
|
|
for dof_name in self.cfg["control"]["stiffness"].keys():
|
||
|
|
if dof_name in joint_name:
|
||
|
|
self.p_gains[i] = self.cfg["control"]["stiffness"][dof_name]
|
||
|
|
self.d_gains[i] = self.cfg["control"]["damping"][dof_name]
|
||
|
|
found = True
|
||
|
|
if not found:
|
||
|
|
self.p_gains[i] = 0.
|
||
|
|
self.d_gains[i] = 0.
|
||
|
|
if self.cfg["control"]["control_type"] in ["P", "V"]:
|
||
|
|
print(f"PD gain of joint {joint_name} were not defined, setting them to zero")
|
||
|
|
|
||
|
|
print(f"p_gains: {self.p_gains}")
|
||
|
|
|
||
|
|
self.commands = np.zeros((1, self.num_commands))
|
||
|
|
self.actions = torch.zeros(12)
|
||
|
|
self.last_actions = torch.zeros(12)
|
||
|
|
self.gravity_vector = np.zeros(3)
|
||
|
|
self.dof_pos = np.zeros(12)
|
||
|
|
self.dof_vel = np.zeros(12)
|
||
|
|
self.body_linear_vel = np.zeros(3)
|
||
|
|
self.body_angular_vel = np.zeros(3)
|
||
|
|
self.joint_pos_target = np.zeros(12)
|
||
|
|
self.joint_vel_target = np.zeros(12)
|
||
|
|
self.torques = np.zeros(12)
|
||
|
|
self.contact_state = np.ones(4)
|
||
|
|
|
||
|
|
self.joint_idxs = self.se.joint_idxs
|
||
|
|
|
||
|
|
self.gait_indices = torch.zeros(self.num_envs, dtype=torch.float)
|
||
|
|
self.clock_inputs = torch.zeros(self.num_envs, 4, dtype=torch.float)
|
||
|
|
|
||
|
|
if "obs_scales" in self.cfg.keys():
|
||
|
|
self.obs_scales = self.cfg["obs_scales"]
|
||
|
|
else:
|
||
|
|
self.obs_scales = self.cfg["normalization"]["obs_scales"]
|
||
|
|
|
||
|
|
self.is_currently_probing = False
|
||
|
|
|
||
|
|
def set_probing(self, is_currently_probing):
|
||
|
|
self.is_currently_probing = is_currently_probing
|
||
|
|
|
||
|
|
def get_obs(self):
|
||
|
|
|
||
|
|
self.gravity_vector = self.se.get_gravity_vector()
|
||
|
|
cmds, reset_timer = self.command_profile.get_command(self.timestep * self.dt, probe=self.is_currently_probing)
|
||
|
|
self.commands[:, :] = cmds[:self.num_commands]
|
||
|
|
if reset_timer:
|
||
|
|
self.reset_gait_indices()
|
||
|
|
#else:
|
||
|
|
# self.commands[:, 0:3] = self.command_profile.get_command(self.timestep * self.dt)[0:3]
|
||
|
|
self.dof_pos = self.se.get_dof_pos()
|
||
|
|
self.dof_vel = self.se.get_dof_vel()
|
||
|
|
self.body_linear_vel = self.se.get_body_linear_vel()
|
||
|
|
self.body_angular_vel = self.se.get_body_angular_vel()
|
||
|
|
|
||
|
|
ob = np.concatenate((self.gravity_vector.reshape(1, -1),
|
||
|
|
self.commands * self.commands_scale,
|
||
|
|
(self.dof_pos - self.default_dof_pos).reshape(1, -1) * self.obs_scales["dof_pos"],
|
||
|
|
self.dof_vel.reshape(1, -1) * self.obs_scales["dof_vel"],
|
||
|
|
torch.clip(self.actions, -self.cfg["normalization"]["clip_actions"],
|
||
|
|
self.cfg["normalization"]["clip_actions"]).cpu().detach().numpy().reshape(1, -1)
|
||
|
|
), axis=1)
|
||
|
|
|
||
|
|
if self.cfg["env"]["observe_two_prev_actions"]:
|
||
|
|
ob = np.concatenate((ob,
|
||
|
|
self.last_actions.cpu().detach().numpy().reshape(1, -1)), axis=1)
|
||
|
|
|
||
|
|
if self.cfg["env"]["observe_clock_inputs"]:
|
||
|
|
ob = np.concatenate((ob,
|
||
|
|
self.clock_inputs), axis=1)
|
||
|
|
# print(self.clock_inputs)
|
||
|
|
|
||
|
|
if self.cfg["env"]["observe_vel"]:
|
||
|
|
ob = np.concatenate(
|
||
|
|
(self.body_linear_vel.reshape(1, -1) * self.obs_scales["lin_vel"],
|
||
|
|
self.body_angular_vel.reshape(1, -1) * self.obs_scales["ang_vel"],
|
||
|
|
ob), axis=1)
|
||
|
|
|
||
|
|
if self.cfg["env"]["observe_only_lin_vel"]:
|
||
|
|
ob = np.concatenate(
|
||
|
|
(self.body_linear_vel.reshape(1, -1) * self.obs_scales["lin_vel"],
|
||
|
|
ob), axis=1)
|
||
|
|
|
||
|
|
if self.cfg["env"]["observe_yaw"]:
|
||
|
|
heading = self.se.get_yaw()
|
||
|
|
ob = np.concatenate((ob, heading.reshape(1, -1)), axis=-1)
|
||
|
|
|
||
|
|
self.contact_state = self.se.get_contact_state()
|
||
|
|
if "observe_contact_states" in self.cfg["env"].keys() and self.cfg["env"]["observe_contact_states"]:
|
||
|
|
ob = np.concatenate((ob, self.contact_state.reshape(1, -1)), axis=-1)
|
||
|
|
|
||
|
|
if "terrain" in self.cfg.keys() and self.cfg["terrain"]["measure_heights"]:
|
||
|
|
robot_height = 0.25
|
||
|
|
self.measured_heights = np.zeros(
|
||
|
|
(len(self.cfg["terrain"]["measured_points_x"]), len(self.cfg["terrain"]["measured_points_y"]))).reshape(
|
||
|
|
1, -1)
|
||
|
|
heights = np.clip(robot_height - 0.5 - self.measured_heights, -1, 1.) * self.obs_scales["height_measurements"]
|
||
|
|
ob = np.concatenate((ob, heights), axis=1)
|
||
|
|
|
||
|
|
|
||
|
|
return torch.tensor(ob, device=self.device).float()
|
||
|
|
|
||
|
|
def get_privileged_observations(self):
|
||
|
|
return None
|
||
|
|
|
||
|
|
def publish_action(self, action, hard_reset=False):
|
||
|
|
|
||
|
|
command_for_robot = pd_tau_targets_lcmt()
|
||
|
|
self.joint_pos_target = \
|
||
|
|
(action[0, :12].detach().cpu().numpy() * self.cfg["control"]["action_scale"]).flatten()
|
||
|
|
self.joint_pos_target[[0, 3, 6, 9]] *= self.cfg["control"]["hip_scale_reduction"]
|
||
|
|
# self.joint_pos_target[[0, 3, 6, 9]] *= -1
|
||
|
|
self.joint_pos_target = self.joint_pos_target
|
||
|
|
self.joint_pos_target += self.default_dof_pos
|
||
|
|
joint_pos_target = self.joint_pos_target[self.joint_idxs]
|
||
|
|
self.joint_vel_target = np.zeros(12)
|
||
|
|
# print(f'cjp {self.joint_pos_target}')
|
||
|
|
|
||
|
|
command_for_robot.q_des = joint_pos_target
|
||
|
|
command_for_robot.qd_des = self.joint_vel_target
|
||
|
|
command_for_robot.kp = self.p_gains
|
||
|
|
command_for_robot.kd = self.d_gains
|
||
|
|
command_for_robot.tau_ff = np.zeros(12)
|
||
|
|
command_for_robot.se_contactState = np.zeros(4)
|
||
|
|
command_for_robot.timestamp_us = int(time.time() * 10 ** 6)
|
||
|
|
command_for_robot.id = 0
|
||
|
|
|
||
|
|
if hard_reset:
|
||
|
|
command_for_robot.id = -1
|
||
|
|
|
||
|
|
|
||
|
|
self.torques = (self.joint_pos_target - self.dof_pos) * self.p_gains + (self.joint_vel_target - self.dof_vel) * self.d_gains
|
||
|
|
|
||
|
|
lc.publish("pd_plustau_targets", command_for_robot.encode())
|
||
|
|
|
||
|
|
def reset(self):
|
||
|
|
self.actions = torch.zeros(12)
|
||
|
|
self.time = time.time()
|
||
|
|
self.timestep = 0
|
||
|
|
return self.get_obs()
|
||
|
|
|
||
|
|
def reset_gait_indices(self):
|
||
|
|
self.gait_indices = torch.zeros(self.num_envs, dtype=torch.float)
|
||
|
|
|
||
|
|
def step(self, actions, hard_reset=False):
|
||
|
|
clip_actions = self.cfg["normalization"]["clip_actions"]
|
||
|
|
self.last_actions = self.actions[:]
|
||
|
|
self.actions = torch.clip(actions[0:1, :], -clip_actions, clip_actions)
|
||
|
|
self.publish_action(self.actions, hard_reset=hard_reset)
|
||
|
|
time.sleep(max(self.dt - (time.time() - self.time), 0))
|
||
|
|
if self.timestep % 100 == 0: print(f'frq: {1 / (time.time() - self.time)} Hz');
|
||
|
|
self.time = time.time()
|
||
|
|
obs = self.get_obs()
|
||
|
|
|
||
|
|
# clock accounting
|
||
|
|
frequencies = self.commands[:, 4]
|
||
|
|
phases = self.commands[:, 5]
|
||
|
|
offsets = self.commands[:, 6]
|
||
|
|
if self.num_commands == 8:
|
||
|
|
bounds = 0
|
||
|
|
durations = self.commands[:, 7]
|
||
|
|
else:
|
||
|
|
bounds = self.commands[:, 7]
|
||
|
|
durations = self.commands[:, 8]
|
||
|
|
self.gait_indices = torch.remainder(self.gait_indices + self.dt * frequencies, 1.0)
|
||
|
|
|
||
|
|
if "pacing_offset" in self.cfg["commands"] and self.cfg["commands"]["pacing_offset"]:
|
||
|
|
self.foot_indices = [self.gait_indices + phases + offsets + bounds,
|
||
|
|
self.gait_indices + bounds,
|
||
|
|
self.gait_indices + offsets,
|
||
|
|
self.gait_indices + phases]
|
||
|
|
else:
|
||
|
|
self.foot_indices = [self.gait_indices + phases + offsets + bounds,
|
||
|
|
self.gait_indices + offsets,
|
||
|
|
self.gait_indices + bounds,
|
||
|
|
self.gait_indices + phases]
|
||
|
|
self.clock_inputs[:, 0] = torch.sin(2 * np.pi * self.foot_indices[0])
|
||
|
|
self.clock_inputs[:, 1] = torch.sin(2 * np.pi * self.foot_indices[1])
|
||
|
|
self.clock_inputs[:, 2] = torch.sin(2 * np.pi * self.foot_indices[2])
|
||
|
|
self.clock_inputs[:, 3] = torch.sin(2 * np.pi * self.foot_indices[3])
|
||
|
|
|
||
|
|
|
||
|
|
images = {'front': self.se.get_camera_front(),
|
||
|
|
'bottom': self.se.get_camera_bottom(),
|
||
|
|
'rear': self.se.get_camera_rear(),
|
||
|
|
'left': self.se.get_camera_left(),
|
||
|
|
'right': self.se.get_camera_right()
|
||
|
|
}
|
||
|
|
downscale_factor = 2
|
||
|
|
temporal_downscale = 3
|
||
|
|
|
||
|
|
for k, v in images.items():
|
||
|
|
if images[k] is not None:
|
||
|
|
images[k] = cv2.resize(images[k], dsize=(images[k].shape[0]//downscale_factor, images[k].shape[1]//downscale_factor), interpolation=cv2.INTER_CUBIC)
|
||
|
|
if self.timestep % temporal_downscale != 0:
|
||
|
|
images[k] = None
|
||
|
|
#print(self.commands)
|
||
|
|
|
||
|
|
infos = {"joint_pos": self.dof_pos[np.newaxis, :],
|
||
|
|
"joint_vel": self.dof_vel[np.newaxis, :],
|
||
|
|
"joint_pos_target": self.joint_pos_target[np.newaxis, :],
|
||
|
|
"joint_vel_target": self.joint_vel_target[np.newaxis, :],
|
||
|
|
"body_linear_vel": self.body_linear_vel[np.newaxis, :],
|
||
|
|
"body_angular_vel": self.body_angular_vel[np.newaxis, :],
|
||
|
|
"contact_state": self.contact_state[np.newaxis, :],
|
||
|
|
"clock_inputs": self.clock_inputs[np.newaxis, :],
|
||
|
|
"body_linear_vel_cmd": self.commands[:, 0:2],
|
||
|
|
"body_angular_vel_cmd": self.commands[:, 2:],
|
||
|
|
"privileged_obs": None,
|
||
|
|
"camera_image_front": images['front'],
|
||
|
|
"camera_image_bottom": images['bottom'],
|
||
|
|
"camera_image_rear": images['rear'],
|
||
|
|
"camera_image_left": images['left'],
|
||
|
|
"camera_image_right": images['right'],
|
||
|
|
}
|
||
|
|
|
||
|
|
self.timestep += 1
|
||
|
|
return obs, None, None, infos
|