Many modifs.
This commit is contained in:
jogima-cyber
parent
2d6ca96367
commit
56c62ae370
|
|
@ -0,0 +1,239 @@
|
|||
<mujoco model="go2">
|
||||
<compiler angle="radian" meshdir="assets" autolimits="true"/>
|
||||
|
||||
<option cone="elliptic" impratio="100"/>
|
||||
|
||||
<default>
|
||||
<default class="go2">
|
||||
<geom friction="0.6" margin="0.001" condim="1"/>
|
||||
<joint axis="0 1 0" damping="0" armature="0.01" frictionloss="0.2"/>
|
||||
<motor ctrlrange="-23.7 23.7"/>
|
||||
<default class="abduction">
|
||||
<joint axis="1 0 0" range="-1.0472 1.0472"/>
|
||||
</default>
|
||||
<default class="hip">
|
||||
<default class="front_hip">
|
||||
<joint range="-1.5708 3.4907"/>
|
||||
</default>
|
||||
<default class="back_hip">
|
||||
<joint range="-0.5236 4.5379"/>
|
||||
</default>
|
||||
</default>
|
||||
<default class="knee">
|
||||
<joint range="-2.7227 -0.83776"/>
|
||||
<motor ctrlrange="-45.43 45.43"/>
|
||||
</default>
|
||||
<default class="visual">
|
||||
<geom type="mesh" contype="0" conaffinity="0" group="2"/>
|
||||
</default>
|
||||
<default class="collision">
|
||||
<geom group="3"/>
|
||||
<default class="foot">
|
||||
<geom size="0.022" pos="-0.002 0 -0.213" priority="1" solimp="0.015 1 0.031" condim="6"
|
||||
friction="0.8 0.02 0.01"/>
|
||||
</default>
|
||||
</default>
|
||||
</default>
|
||||
</default>
|
||||
|
||||
<asset>
|
||||
<hfield name="terrain" nrow="300" ncol="300" size="5 5 1. 0.1"/>
|
||||
<material name="metal" rgba=".9 .95 .95 1"/>
|
||||
<material name="black" rgba="0 0 0 1"/>
|
||||
<material name="white" rgba="1 1 1 1"/>
|
||||
<material name="gray" rgba="0.671705 0.692426 0.774270 1"/>
|
||||
<texture type="skybox" builtin="gradient" rgb1="0.3 0.5 0.7" rgb2="0 0 0" width="512" height="3072"/>
|
||||
<texture type="2d" name="groundplane" builtin="checker" mark="edge" rgb1="0.2 0.3 0.4" rgb2="0.1 0.2 0.3" markrgb="0.8 0.8 0.8" width="300" height="300"/>
|
||||
<material name="groundplane" texture="groundplane" texuniform="true" texrepeat="5 5" reflectance="0.2"/>
|
||||
|
||||
<mesh file="base_0.obj"/>
|
||||
<mesh file="base_1.obj"/>
|
||||
<mesh file="base_2.obj"/>
|
||||
<mesh file="base_3.obj"/>
|
||||
<mesh file="base_4.obj"/>
|
||||
<mesh file="hip_0.obj"/>
|
||||
<mesh file="hip_1.obj"/>
|
||||
<mesh file="thigh_0.obj"/>
|
||||
<mesh file="thigh_1.obj"/>
|
||||
<mesh file="thigh_mirror_0.obj"/>
|
||||
<mesh file="thigh_mirror_1.obj"/>
|
||||
<mesh file="calf_0.obj"/>
|
||||
<mesh file="calf_1.obj"/>
|
||||
<mesh file="calf_mirror_0.obj"/>
|
||||
<mesh file="calf_mirror_1.obj"/>
|
||||
<mesh file="foot.obj"/>
|
||||
</asset>
|
||||
|
||||
<worldbody>
|
||||
<geom name="heightfield" pos="0.0 0.0 -0.1" type="hfield" hfield="terrain"/>
|
||||
<geom name="floor" size="0 0 0.05" type="plane" material="groundplane"/>
|
||||
<body name="base" pos="0 0 0.445" childclass="go2">
|
||||
<light pos="0.0 0.0 2.5" />
|
||||
<camera name="front_camera" mode="fixed" pos="0.321618 0.033305 0.081622" euler="1.57079632679 0. 0." resolution="640 480" fovy="45."/>
|
||||
<inertial pos="0.021112 0 -0.005366" quat="-0.000543471 0.713435 -0.00173769 0.700719" mass="6.921"
|
||||
diaginertia="0.107027 0.0980771 0.0244531"/>
|
||||
<freejoint/>
|
||||
<geom mesh="base_0" material="black" class="visual"/>
|
||||
<geom mesh="base_1" material="black" class="visual"/>
|
||||
<geom mesh="base_2" material="black" class="visual"/>
|
||||
<geom mesh="base_3" material="white" class="visual"/>
|
||||
<geom mesh="base_4" material="gray" class="visual"/>
|
||||
<geom size="0.1881 0.04675 0.057" type="box" class="collision"/>
|
||||
<geom size="0.05 0.045" pos="0.285 0 0.01" type="cylinder" class="collision"/>
|
||||
<geom size="0.047" pos="0.293 0 -0.06" class="collision"/>
|
||||
<site name="imu" pos="-0.02557 0 0.04232"/>
|
||||
|
||||
<body name="FR_hip" pos="0.1934 -0.0465 0">
|
||||
<inertial pos="-0.0054 -0.00194 -0.000105" quat="0.498237 0.505462 0.499245 0.497014" mass="0.678"
|
||||
diaginertia="0.00088403 0.000596003 0.000479967"/>
|
||||
<joint name="FR_hip_joint" class="abduction"/>
|
||||
<geom mesh="hip_0" material="metal" class="visual" quat="4.63268e-05 1 0 0"/>
|
||||
<geom mesh="hip_1" material="gray" class="visual" quat="4.63268e-05 1 0 0"/>
|
||||
<geom size="0.046 0.02" pos="0 -0.08 0" quat="0.707107 0.707107 0 0" type="cylinder" class="collision"/>
|
||||
<body name="FR_thigh" pos="0 -0.0955 0">
|
||||
<inertial pos="-0.00374 0.0223 -0.0327" quat="0.551623 -0.0200632 0.0847635 0.829533" mass="1.152"
|
||||
diaginertia="0.00594973 0.00584149 0.000878787"/>
|
||||
<joint name="FR_thigh_joint" class="front_hip"/>
|
||||
<geom mesh="thigh_mirror_0" material="metal" class="visual"/>
|
||||
<geom mesh="thigh_mirror_1" material="gray" class="visual"/>
|
||||
<geom size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0" type="box" class="collision"/>
|
||||
<body name="FR_calf" pos="0 0 -0.213">
|
||||
<inertial pos="0.00629595 0.000622121 -0.141417" quat="0.703508 -0.00450087 0.00154099 0.710672"
|
||||
mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05"/>
|
||||
<joint name="FR_calf_joint" class="knee"/>
|
||||
<geom mesh="calf_mirror_0" material="gray" class="visual"/>
|
||||
<geom mesh="calf_mirror_1" material="black" class="visual"/>
|
||||
<geom size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0" type="cylinder" class="collision"/>
|
||||
<geom size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" type="cylinder" class="collision"/>
|
||||
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black"/>
|
||||
<geom name="FR" class="foot"/>
|
||||
<site name="FR_foot" pos="0 0 -0.213" size="0.02"/>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
|
||||
<body name="FL_hip" pos="0.1934 0.0465 0">
|
||||
<inertial pos="-0.0054 0.00194 -0.000105" quat="0.497014 0.499245 0.505462 0.498237" mass="0.678"
|
||||
diaginertia="0.00088403 0.000596003 0.000479967"/>
|
||||
<joint name="FL_hip_joint" class="abduction"/>
|
||||
<geom mesh="hip_0" material="metal" class="visual"/>
|
||||
<geom mesh="hip_1" material="gray" class="visual"/>
|
||||
<geom size="0.046 0.02" pos="0 0.08 0" quat="1 1 0 0" type="cylinder" class="collision"/>
|
||||
<body name="FL_thigh" pos="0 0.0955 0">
|
||||
<inertial pos="-0.00374 -0.0223 -0.0327" quat="0.829533 0.0847635 -0.0200632 0.551623" mass="1.152"
|
||||
diaginertia="0.00594973 0.00584149 0.000878787"/>
|
||||
<joint name="FL_thigh_joint" class="front_hip"/>
|
||||
<geom mesh="thigh_0" material="metal" class="visual"/>
|
||||
<geom mesh="thigh_1" material="gray" class="visual"/>
|
||||
<geom size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0" type="box" class="collision"/>
|
||||
<body name="FL_calf" pos="0 0 -0.213">
|
||||
<inertial pos="0.00629595 -0.000622121 -0.141417" quat="0.710672 0.00154099 -0.00450087 0.703508"
|
||||
mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05"/>
|
||||
<joint name="FL_calf_joint" class="knee"/>
|
||||
<geom mesh="calf_0" material="gray" class="visual"/>
|
||||
<geom mesh="calf_1" material="black" class="visual"/>
|
||||
<geom size="0.012 0.06" pos="0.008 0 -0.06" quat="0.994493 0 -0.104807 0" type="cylinder" class="collision"/>
|
||||
<geom size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" type="cylinder" class="collision"/>
|
||||
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black"/>
|
||||
<geom name="FL" class="foot"/>
|
||||
<site name="FL_foot" pos="0 0 -0.213" size="0.02"/>
|
||||
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
|
||||
<body name="RR_hip" pos="-0.1934 -0.0465 0">
|
||||
<inertial pos="0.0054 -0.00194 -0.000105" quat="0.499245 0.497014 0.498237 0.505462" mass="0.678"
|
||||
diaginertia="0.00088403 0.000596003 0.000479967"/>
|
||||
<joint name="RR_hip_joint" class="abduction"/>
|
||||
<geom mesh="hip_0" material="metal" class="visual" quat="2.14617e-09 4.63268e-05 4.63268e-05 -1"/>
|
||||
<geom mesh="hip_1" material="gray" class="visual" quat="2.14617e-09 4.63268e-05 4.63268e-05 -1"/>
|
||||
<geom size="0.046 0.02" pos="0 -0.08 0" quat="0.707107 0.707107 0 0" type="cylinder" class="collision"/>
|
||||
<body name="RR_thigh" pos="0 -0.0955 0">
|
||||
<inertial pos="-0.00374 0.0223 -0.0327" quat="0.551623 -0.0200632 0.0847635 0.829533" mass="1.152"
|
||||
diaginertia="0.00594973 0.00584149 0.000878787"/>
|
||||
<joint name="RR_thigh_joint" class="back_hip"/>
|
||||
<geom mesh="thigh_mirror_0" material="metal" class="visual"/>
|
||||
<geom mesh="thigh_mirror_1" material="gray" class="visual"/>
|
||||
<geom size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0" type="box" class="collision"/>
|
||||
<body name="RR_calf" pos="0 0 -0.213">
|
||||
<inertial pos="0.00629595 0.000622121 -0.141417" quat="0.703508 -0.00450087 0.00154099 0.710672"
|
||||
mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05"/>
|
||||
<joint name="RR_calf_joint" class="knee"/>
|
||||
<geom mesh="calf_mirror_0" material="gray" class="visual"/>
|
||||
<geom mesh="calf_mirror_1" material="black" class="visual"/>
|
||||
<geom size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0" type="cylinder" class="collision"/>
|
||||
<geom size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" type="cylinder" class="collision"/>
|
||||
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black"/>
|
||||
<geom name="RR" class="foot"/>
|
||||
<site name="RR_foot" pos="0 0 -0.213" size="0.02"/>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
|
||||
<body name="RL_hip" pos="-0.1934 0.0465 0">
|
||||
<inertial pos="0.0054 0.00194 -0.000105" quat="0.505462 0.498237 0.497014 0.499245" mass="0.678"
|
||||
diaginertia="0.00088403 0.000596003 0.000479967"/>
|
||||
<joint name="RL_hip_joint" class="abduction"/>
|
||||
<geom mesh="hip_0" material="metal" class="visual" quat="4.63268e-05 0 1 0"/>
|
||||
<geom mesh="hip_1" material="gray" class="visual" quat="4.63268e-05 0 1 0"/>
|
||||
<geom size="0.046 0.02" pos="0 0.08 0" quat="0.707107 0.707107 0 0" type="cylinder" class="collision"/>
|
||||
<body name="RL_thigh" pos="0 0.0955 0">
|
||||
<inertial pos="-0.00374 -0.0223 -0.0327" quat="0.829533 0.0847635 -0.0200632 0.551623" mass="1.152"
|
||||
diaginertia="0.00594973 0.00584149 0.000878787"/>
|
||||
<joint name="RL_thigh_joint" class="back_hip"/>
|
||||
<geom mesh="thigh_0" material="metal" class="visual"/>
|
||||
<geom mesh="thigh_1" material="gray" class="visual"/>
|
||||
<geom size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0" type="box" class="collision"/>
|
||||
<body name="RL_calf" pos="0 0 -0.213">
|
||||
<inertial pos="0.00629595 -0.000622121 -0.141417" quat="0.710672 0.00154099 -0.00450087 0.703508"
|
||||
mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05"/>
|
||||
<joint name="RL_calf_joint" class="knee"/>
|
||||
<geom mesh="calf_0" material="gray" class="visual"/>
|
||||
<geom mesh="calf_1" material="black" class="visual"/>
|
||||
<geom size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0" type="cylinder" class="collision"/>
|
||||
<geom size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" type="cylinder" class="collision"/>
|
||||
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black"/>
|
||||
<geom name="RL" class="foot"/>
|
||||
<site name="RL_foot" pos="0 0 -0.213" size="0.02"/>
|
||||
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
|
||||
|
||||
|
||||
</body>
|
||||
</worldbody>
|
||||
|
||||
<actuator>
|
||||
<motor class="abduction" name="FR_hip" joint="FR_hip_joint"/>
|
||||
<motor class="hip" name="FR_thigh" joint="FR_thigh_joint"/>
|
||||
<motor class="knee" name="FR_calf" joint="FR_calf_joint"/>
|
||||
<motor class="abduction" name="FL_hip" joint="FL_hip_joint"/>
|
||||
<motor class="hip" name="FL_thigh" joint="FL_thigh_joint"/>
|
||||
<motor class="knee" name="FL_calf" joint="FL_calf_joint"/>
|
||||
<motor class="abduction" name="RR_hip" joint="RR_hip_joint"/>
|
||||
<motor class="hip" name="RR_thigh" joint="RR_thigh_joint"/>
|
||||
<motor class="knee" name="RR_calf" joint="RR_calf_joint"/>
|
||||
<motor class="abduction" name="RL_hip" joint="RL_hip_joint"/>
|
||||
<motor class="hip" name="RL_thigh" joint="RL_thigh_joint"/>
|
||||
<motor class="knee" name="RL_calf" joint="RL_calf_joint"/>
|
||||
</actuator>
|
||||
|
||||
<sensor>
|
||||
<!-- IMU -->
|
||||
<accelerometer name="imu_accel" site="imu"/>
|
||||
<gyro name="imu_gyro" site="imu"/>
|
||||
<!-- Contact sensors for the feet -->
|
||||
<touch name="FR_foot_contact" site="FR_foot"/>
|
||||
<touch name="FL_foot_contact" site="RL_foot"/>
|
||||
<touch name="RR_foot_contact" site="RR_foot"/>
|
||||
<touch name="RL_foot_contact" site="RL_foot"/>
|
||||
</sensor>
|
||||
|
||||
<keyframe>
|
||||
<key name="home" qpos="0 0 0.27 1 0 0 0 0 0.9 -1.8 0 0.9 -1.8 0 0.9 -1.8 0 0.9 -1.8"
|
||||
ctrl="0 0.9 -1.8 0 0.9 -1.8 0 0.9 -1.8 0 0.9 -1.8"/>
|
||||
</keyframe>
|
||||
</mujoco>
|
||||
|
|
@ -666,71 +666,6 @@ Stephen Brawner (brawner@gmail.com)
|
|||
<axis xyz="0 0 0"/>
|
||||
</joint>
|
||||
|
||||
<joint name="vicon_joint" type="fixed">
|
||||
<parent link="base_link"/>
|
||||
<child link="vicon_link"/>
|
||||
<origin rpy="-0.037935 -0.066371 -2.668897" xyz="0.022944 -0.000007 0.058458"/>
|
||||
</joint>
|
||||
|
||||
<link name="vicon_link">
|
||||
<visual>
|
||||
<origin rpy="0 0 0" xyz="0 0 0"/>
|
||||
<geometry>
|
||||
<box size="0.001 0.001 0.001"/>
|
||||
</geometry>
|
||||
<!-- <material name="red"/> -->
|
||||
</visual>
|
||||
</link>
|
||||
|
||||
|
||||
<joint name="infra1_joint" type="fixed">
|
||||
<parent link="base_link"/>
|
||||
<child link="infra1_link"/>
|
||||
<origin rpy="-1.641669 -0.014332 -1.617155" xyz="0.321618 0.033305 0.081622"/>
|
||||
</joint>
|
||||
|
||||
<link name="infra1_link">
|
||||
<visual>
|
||||
<origin rpy="0 0 0" xyz="0 0 0"/>
|
||||
<geometry>
|
||||
<box size="0.001 0.001 0.001"/>
|
||||
</geometry>
|
||||
<!-- <material name="red"/> -->
|
||||
</visual>
|
||||
</link>
|
||||
|
||||
<joint name="infra2_joint" type="fixed">
|
||||
<parent link="base_link"/>
|
||||
<child link="infra1_link"/>
|
||||
<origin rpy="-1.640637 -0.014093 -1.611177" xyz="0.317969 -0.061919 0.082773"/>
|
||||
</joint>
|
||||
|
||||
<link name="infra2_link">
|
||||
<visual>
|
||||
<origin rpy="0 0 0" xyz="0 0 0"/>
|
||||
<geometry>
|
||||
<box size="0.001 0.001 0.001"/>
|
||||
</geometry>
|
||||
<!-- <material name="red"/> -->
|
||||
</visual>
|
||||
</link>
|
||||
|
||||
<joint name="color_joint" type="fixed">
|
||||
<parent link="base_link"/>
|
||||
<child link="color_link"/>
|
||||
<origin rpy="-1.640491 -0.011125 -1.610364" xyz="0.319415 -0.025790 0.082311"/>
|
||||
</joint>
|
||||
|
||||
<link name="color_link">
|
||||
<visual>
|
||||
<origin rpy="0 0 0" xyz="0 0 0"/>
|
||||
<geometry>
|
||||
<box size="0.001 0.001 0.001"/>
|
||||
</geometry>
|
||||
<!-- <material name="red"/> -->
|
||||
</visual>
|
||||
</link>
|
||||
|
||||
<link name="utlidar">
|
||||
<inertial>
|
||||
<origin rpy="0 0 0" xyz="0 0 0"/>
|
||||
|
|
|
|||
|
|
@ -235,6 +235,7 @@ class CaTAgent:
|
|||
self.gravity_vector = np.zeros(3)
|
||||
self.dof_pos = np.zeros(12)
|
||||
self.dof_vel = np.zeros(12)
|
||||
self.pos = np.zeros(3)
|
||||
self.body_linear_vel = np.zeros(3)
|
||||
self.body_angular_vel = np.zeros(3)
|
||||
self.joint_pos_target = np.zeros(12)
|
||||
|
|
@ -263,9 +264,11 @@ class CaTAgent:
|
|||
self.prev_dof_vel = self.dof_vel.copy()
|
||||
self.dof_vel = np.array(joint_state['dq'])[self.unitree_to_policy_map[:, 1]]
|
||||
self.body_angular_vel = self.robot.getIMU()["gyro"]
|
||||
self.body_linear_vel = self.robot.getLinVel()
|
||||
|
||||
try:
|
||||
self.foot_contact_forces_mag = self.robot.getFootContact()
|
||||
self.body_linear_vel = self.robot.getLinVel()
|
||||
self.pos, _ = self.robot.getPose()
|
||||
except:
|
||||
pass
|
||||
|
||||
|
|
@ -327,7 +330,7 @@ class CaTAgent:
|
|||
# if self.timestep % 100 == 0:
|
||||
# print(f"frq: {1 / (time.time() - self.time)} Hz")
|
||||
self.time = time.time()
|
||||
obs = self.get_obs()
|
||||
#obs = self.get_obs()
|
||||
|
||||
joint_acc = np.abs(self.prev_dof_vel - self.dof_vel) / self.dt
|
||||
if self.prev_joint_acc is None:
|
||||
|
|
@ -345,6 +348,7 @@ class CaTAgent:
|
|||
"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, :],
|
||||
"body_pos": self.pos[np.newaxis, :].copy(),
|
||||
"contact_state": self.contact_state[np.newaxis, :],
|
||||
"body_linear_vel_cmd": self.commands[np.newaxis, 0:2],
|
||||
"body_angular_vel_cmd": self.commands[np.newaxis, 2:],
|
||||
|
|
@ -356,4 +360,4 @@ class CaTAgent:
|
|||
}
|
||||
|
||||
self.timestep += 1
|
||||
return obs, None, None, infos
|
||||
return None, None, None, infos
|
||||
|
|
|
|||
|
|
@ -0,0 +1,364 @@
|
|||
import os
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
import onnxruntime
|
||||
|
||||
class RunningMeanStd(nn.Module):
|
||||
def __init__(self, shape = (), epsilon=1e-08):
|
||||
super(RunningMeanStd, self).__init__()
|
||||
self.register_buffer("running_mean", torch.zeros(shape))
|
||||
self.register_buffer("running_var", torch.ones(shape))
|
||||
self.register_buffer("count", torch.ones(()))
|
||||
|
||||
self.epsilon = epsilon
|
||||
|
||||
def forward(self, obs, update = True):
|
||||
if update:
|
||||
self.update(obs)
|
||||
|
||||
return (obs - self.running_mean) / torch.sqrt(self.running_var + self.epsilon)
|
||||
|
||||
def update(self, x):
|
||||
"""Updates the mean, var and count from a batch of samples."""
|
||||
batch_mean = torch.mean(x, dim=0)
|
||||
batch_var = torch.var(x, correction=0, dim=0)
|
||||
batch_count = x.shape[0]
|
||||
self.update_from_moments(batch_mean, batch_var, batch_count)
|
||||
|
||||
def update_from_moments(self, batch_mean, batch_var, batch_count):
|
||||
"""Updates from batch mean, variance and count moments."""
|
||||
self.running_mean, self.running_var, self.count = update_mean_var_count_from_moments(
|
||||
self.running_mean, self.running_var, self.count, batch_mean, batch_var, batch_count
|
||||
)
|
||||
|
||||
def update_mean_var_count_from_moments(
|
||||
mean, var, count, batch_mean, batch_var, batch_count
|
||||
):
|
||||
"""Updates the mean, var and count using the previous mean, var, count and batch values."""
|
||||
delta = batch_mean - mean
|
||||
tot_count = count + batch_count
|
||||
|
||||
new_mean = mean + delta * batch_count / tot_count
|
||||
m_a = var * count
|
||||
m_b = batch_var * batch_count
|
||||
M2 = m_a + m_b + torch.square(delta) * count * batch_count / tot_count
|
||||
new_var = M2 / tot_count
|
||||
new_count = tot_count
|
||||
|
||||
return new_mean, new_var, new_count
|
||||
|
||||
def layer_init(layer, std=np.sqrt(2), bias_const=0.0):
|
||||
torch.nn.init.orthogonal_(layer.weight, std)
|
||||
torch.nn.init.constant_(layer.bias, bias_const)
|
||||
return layer
|
||||
|
||||
class Agent(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.critic = nn.Sequential(
|
||||
layer_init(nn.Linear(45 + 144, 512)),
|
||||
nn.ELU(),
|
||||
layer_init(nn.Linear(512, 256)),
|
||||
nn.ELU(),
|
||||
layer_init(nn.Linear(256, 128)),
|
||||
nn.ELU(),
|
||||
layer_init(nn.Linear(128, 1), std=1.0),
|
||||
)
|
||||
self.actor_mean = nn.Sequential(
|
||||
layer_init(nn.Linear(45 + 144, 512)),
|
||||
nn.ELU(),
|
||||
layer_init(nn.Linear(512, 256)),
|
||||
nn.ELU(),
|
||||
layer_init(nn.Linear(256, 128)),
|
||||
nn.ELU(),
|
||||
layer_init(nn.Linear(128, 12), std=0.01),
|
||||
)
|
||||
self.actor_logstd = nn.Parameter(torch.zeros(1, 12))
|
||||
|
||||
self.obs_rms = RunningMeanStd(shape = (45 + 144,))
|
||||
self.value_rms = RunningMeanStd(shape = ())
|
||||
|
||||
def get_value(self, x):
|
||||
return self.critic(x)
|
||||
|
||||
def get_action(self, x):
|
||||
action_mean = self.actor_mean(x)
|
||||
return action_mean
|
||||
|
||||
def forward(self, x):
|
||||
action_mean = self.actor_mean(self.obs_rms(x, update = False))
|
||||
return action_mean
|
||||
|
||||
class Policy:
|
||||
def __init__(self, checkpoint_path):
|
||||
self.agent = Agent()
|
||||
actor_sd = torch.load(checkpoint_path, map_location="cpu")
|
||||
self.agent.load_state_dict(actor_sd)
|
||||
onnx_file_name = checkpoint_path.replace(".pt", ".onnx")
|
||||
|
||||
dummy_input = torch.randn(1, 45 + 144)
|
||||
with torch.no_grad():
|
||||
torch_out = self.agent(dummy_input)
|
||||
|
||||
torch.onnx.export(
|
||||
self.agent, # The model being converted
|
||||
dummy_input, # An example input for the model
|
||||
onnx_file_name, # Output file name
|
||||
export_params=True, # Store trained parameter weights inside the model file
|
||||
opset_version=11, # ONNX version (opset) to export to, adjust as needed
|
||||
do_constant_folding=True, # Whether to perform constant folding optimization
|
||||
input_names=['input'], # Name of the input in the ONNX graph (can be customized)
|
||||
output_names=['action'], # Name of the output (assuming get_action and get_value are key outputs)
|
||||
)
|
||||
|
||||
self.ort_session = onnxruntime.InferenceSession(onnx_file_name, providers=["CPUExecutionProvider"])
|
||||
ort_inputs = {'input': dummy_input.numpy()}
|
||||
ort_outs = self.ort_session.run(None, ort_inputs)
|
||||
np.testing.assert_allclose(torch_out.numpy(), ort_outs[0], rtol=1e-03, atol=1e-05)
|
||||
print("Exported model has been tested with ONNXRuntime, and the result looks good!")
|
||||
|
||||
def __call__(self, obs, info):
|
||||
#with torch.no_grad():
|
||||
# action = self.agent.get_action(self.agent.obs_rms(obs.unsqueeze(0), update = False))
|
||||
ort_inputs = {'input': obs[np.newaxis].astype(np.float32)}
|
||||
ort_outs = self.ort_session.run(None, ort_inputs)
|
||||
return ort_outs[0]
|
||||
|
||||
class CommandInterface:
|
||||
def __init__(self, limits=None):
|
||||
self.limits = limits
|
||||
self.x_vel_cmd, self.y_vel_cmd, self.yaw_vel_cmd = 0.0, 0.0, 0.0
|
||||
|
||||
def get_command(self):
|
||||
command = np.zeros((3,))
|
||||
command[0] = self.x_vel_cmd
|
||||
command[1] = self.y_vel_cmd
|
||||
command[2] = self.yaw_vel_cmd
|
||||
return command, False
|
||||
|
||||
|
||||
class CaTAgent:
|
||||
def __init__(self, command_profile, robot):
|
||||
self.robot = robot
|
||||
self.command_profile = command_profile
|
||||
# self.lcm_bridge = LCMBridgeClient(robot_name=self.robot_name)
|
||||
self.sim_dt = 0.001
|
||||
self.decimation = 20
|
||||
self.dt = self.sim_dt * self.decimation
|
||||
self.timestep = 0
|
||||
|
||||
self.device = "cpu"
|
||||
|
||||
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",
|
||||
]
|
||||
|
||||
policy_joint_names = joint_names
|
||||
|
||||
unitree_joint_names = [
|
||||
"FR_hip_joint",
|
||||
"FR_thigh_joint",
|
||||
"FR_calf_joint",
|
||||
"FL_hip_joint",
|
||||
"FL_thigh_joint",
|
||||
"FL_calf_joint",
|
||||
"RR_hip_joint",
|
||||
"RR_thigh_joint",
|
||||
"RR_calf_joint",
|
||||
"RL_hip_joint",
|
||||
"RL_thigh_joint",
|
||||
"RL_calf_joint",
|
||||
]
|
||||
|
||||
policy_to_unitree_map = []
|
||||
unitree_to_policy_map = []
|
||||
for i, policy_joint_name in enumerate(policy_joint_names):
|
||||
id = np.where([name == policy_joint_name for name in unitree_joint_names])[0][0]
|
||||
policy_to_unitree_map.append((i, id))
|
||||
self.policy_to_unitree_map = np.array(policy_to_unitree_map).astype(np.uint32)
|
||||
|
||||
for i, unitree_joint_name in enumerate(unitree_joint_names):
|
||||
id = np.where([name == unitree_joint_name for name in policy_joint_names])[0][0]
|
||||
unitree_to_policy_map.append((i, id))
|
||||
self.unitree_to_policy_map = np.array(unitree_to_policy_map).astype(np.uint32)
|
||||
|
||||
default_joint_angles = {
|
||||
"FL_hip_joint": 0.1,
|
||||
"RL_hip_joint": 0.1,
|
||||
"FR_hip_joint": -0.1,
|
||||
"RR_hip_joint": -0.1,
|
||||
"FL_thigh_joint": 0.8,
|
||||
"RL_thigh_joint": 1.0,
|
||||
"FR_thigh_joint": 0.8,
|
||||
"RR_thigh_joint": 1.0,
|
||||
"FL_calf_joint": -1.5,
|
||||
"RL_calf_joint": -1.5,
|
||||
"FR_calf_joint": -1.5,
|
||||
"RR_calf_joint": -1.5
|
||||
}
|
||||
self.default_dof_pos = np.array(
|
||||
[
|
||||
default_joint_angles[name]
|
||||
for name in joint_names
|
||||
]
|
||||
)
|
||||
|
||||
self.default_dof_pos = self.default_dof_pos
|
||||
|
||||
self.p_gains = np.zeros(12)
|
||||
self.d_gains = np.zeros(12)
|
||||
for i in range(12):
|
||||
self.p_gains[i] = 20.0
|
||||
self.d_gains[i] = 0.5
|
||||
|
||||
print(f"p_gains: {self.p_gains}")
|
||||
|
||||
self.commands = np.zeros(3)
|
||||
self.actions = np.zeros((1, 12))
|
||||
self.last_actions = np.zeros((1,12))
|
||||
self.gravity_vector = np.zeros(3)
|
||||
self.dof_pos = np.zeros(12)
|
||||
self.dof_vel = np.zeros(12)
|
||||
self.pos = np.zeros(3)
|
||||
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.prev_joint_acc = None
|
||||
self.torques = np.zeros(12)
|
||||
self.contact_state = np.ones(4)
|
||||
self.foot_contact_forces_mag = np.zeros(4)
|
||||
self.prev_foot_contact_forces_mag = np.zeros(4)
|
||||
self.test = 0
|
||||
|
||||
def wait_for_state(self):
|
||||
# return self.lcm_bridge.getStates(timeout=2)
|
||||
pass
|
||||
|
||||
def get_obs(self):
|
||||
cmds, reset_timer = self.command_profile.get_command()
|
||||
self.commands[:] = cmds
|
||||
|
||||
# self.state = self.wait_for_state()
|
||||
joint_state = self.robot.getJointStates()
|
||||
if joint_state is not None:
|
||||
self.gravity_vector = self.robot.getGravityInBody()
|
||||
self.prev_dof_pos = self.dof_pos.copy()
|
||||
self.dof_pos = np.array(joint_state['q'])[self.unitree_to_policy_map[:, 1]]
|
||||
self.prev_dof_vel = self.dof_vel.copy()
|
||||
self.dof_vel = np.array(joint_state['dq'])[self.unitree_to_policy_map[:, 1]]
|
||||
self.body_angular_vel = self.robot.getIMU()["gyro"]
|
||||
|
||||
try:
|
||||
self.foot_contact_forces_mag = self.robot.getFootContact()
|
||||
self.body_linear_vel = self.robot.getLinVel()
|
||||
self.pos, _ = self.robot.getPose()
|
||||
except:
|
||||
pass
|
||||
|
||||
ob = np.concatenate(
|
||||
(
|
||||
self.body_angular_vel * 0.25,
|
||||
self.commands * np.array([2.0, 2.0, 0.25]),
|
||||
self.gravity_vector[:, 0],
|
||||
self.dof_pos * 1.0,
|
||||
#((self.dof_pos - self.prev_dof_pos) / self.dt) * 0.05,
|
||||
self.dof_vel * 0.05,
|
||||
self.last_actions[0],
|
||||
np.zeros(144)
|
||||
),
|
||||
axis=0,
|
||||
)
|
||||
|
||||
#return torch.tensor(ob, device=self.device).float()
|
||||
return ob
|
||||
|
||||
def publish_action(self, action, hard_reset=False):
|
||||
# command_for_robot = UnitreeLowCommand()
|
||||
#self.joint_pos_target = (
|
||||
# action[0, :12].detach().cpu().numpy() * 0.25
|
||||
#).flatten()
|
||||
self.joint_pos_target = (
|
||||
action[0, :12] * 0.25
|
||||
).flatten()
|
||||
self.joint_pos_target += self.default_dof_pos
|
||||
self.joint_vel_target = np.zeros(12)
|
||||
# command_for_robot.q_des = self.joint_pos_target
|
||||
# command_for_robot.dq_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)
|
||||
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
|
||||
# self.lcm_bridge.sendCommands(command_for_robot)
|
||||
self.robot.setCommands(self.joint_pos_target[self.policy_to_unitree_map[:, 1]],
|
||||
self.joint_vel_target[self.policy_to_unitree_map[:, 1]],
|
||||
self.p_gains[self.policy_to_unitree_map[:, 1]],
|
||||
self.d_gains[self.policy_to_unitree_map[:, 1]],
|
||||
np.zeros(12))
|
||||
|
||||
def reset(self):
|
||||
self.actions = torch.zeros((1, 12))
|
||||
self.time = time.time()
|
||||
self.timestep = 0
|
||||
return self.get_obs()
|
||||
|
||||
def step(self, actions, hard_reset=False):
|
||||
self.last_actions = self.actions[:]
|
||||
self.actions = 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()
|
||||
|
||||
joint_acc = np.abs(self.prev_dof_vel - self.dof_vel) / self.dt
|
||||
if self.prev_joint_acc is None:
|
||||
self.prev_joint_acc = np.zeros_like(joint_acc)
|
||||
joint_jerk = np.abs(self.prev_joint_acc - joint_acc) / self.dt
|
||||
self.prev_joint_acc = joint_acc.copy()
|
||||
|
||||
foot_contact_rate = np.abs(self.foot_contact_forces_mag - self.prev_foot_contact_forces_mag)
|
||||
self.prev_foot_contact_forces_mag = self.foot_contact_forces_mag.copy()
|
||||
|
||||
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, :],
|
||||
"body_pos": self.pos[np.newaxis, :].copy(),
|
||||
"contact_state": self.contact_state[np.newaxis, :],
|
||||
"body_linear_vel_cmd": self.commands[np.newaxis, 0:2],
|
||||
"body_angular_vel_cmd": self.commands[np.newaxis, 2:],
|
||||
"torques": self.torques,
|
||||
"foot_contact_forces_mag": self.foot_contact_forces_mag.copy(),
|
||||
"joint_acc": joint_acc[np.newaxis, :],
|
||||
"joint_jerk": joint_jerk[np.newaxis, :],
|
||||
"foot_contact_rate": foot_contact_rate[np.newaxis, :],
|
||||
}
|
||||
|
||||
self.timestep += 1
|
||||
return None, None, None, infos
|
||||
|
|
@ -0,0 +1,361 @@
|
|||
import os
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
import onnxruntime
|
||||
|
||||
class RunningMeanStd(nn.Module):
|
||||
def __init__(self, shape = (), epsilon=1e-08):
|
||||
super(RunningMeanStd, self).__init__()
|
||||
self.register_buffer("running_mean", torch.zeros(shape))
|
||||
self.register_buffer("running_var", torch.ones(shape))
|
||||
self.register_buffer("count", torch.ones(()))
|
||||
|
||||
self.epsilon = epsilon
|
||||
|
||||
def forward(self, obs, update = True):
|
||||
if update:
|
||||
self.update(obs)
|
||||
|
||||
return (obs - self.running_mean) / torch.sqrt(self.running_var + self.epsilon)
|
||||
|
||||
def update(self, x):
|
||||
"""Updates the mean, var and count from a batch of samples."""
|
||||
batch_mean = torch.mean(x, dim=0)
|
||||
batch_var = torch.var(x, correction=0, dim=0)
|
||||
batch_count = x.shape[0]
|
||||
self.update_from_moments(batch_mean, batch_var, batch_count)
|
||||
|
||||
def update_from_moments(self, batch_mean, batch_var, batch_count):
|
||||
"""Updates from batch mean, variance and count moments."""
|
||||
self.running_mean, self.running_var, self.count = update_mean_var_count_from_moments(
|
||||
self.running_mean, self.running_var, self.count, batch_mean, batch_var, batch_count
|
||||
)
|
||||
|
||||
def update_mean_var_count_from_moments(
|
||||
mean, var, count, batch_mean, batch_var, batch_count
|
||||
):
|
||||
"""Updates the mean, var and count using the previous mean, var, count and batch values."""
|
||||
delta = batch_mean - mean
|
||||
tot_count = count + batch_count
|
||||
|
||||
new_mean = mean + delta * batch_count / tot_count
|
||||
m_a = var * count
|
||||
m_b = batch_var * batch_count
|
||||
M2 = m_a + m_b + torch.square(delta) * count * batch_count / tot_count
|
||||
new_var = M2 / tot_count
|
||||
new_count = tot_count
|
||||
|
||||
return new_mean, new_var, new_count
|
||||
|
||||
def layer_init(layer, std=np.sqrt(2), bias_const=0.0):
|
||||
torch.nn.init.orthogonal_(layer.weight, std)
|
||||
torch.nn.init.constant_(layer.bias, bias_const)
|
||||
return layer
|
||||
|
||||
class Agent(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.actor_memory = nn.GRU(45, 256, batch_first=True)
|
||||
self.critic_memory = nn.GRU(45, 256, batch_first=True)
|
||||
self.critic = nn.Sequential(
|
||||
layer_init(nn.Linear(45 + 256, 512)),
|
||||
nn.ELU(),
|
||||
layer_init(nn.Linear(512, 256)),
|
||||
nn.ELU(),
|
||||
layer_init(nn.Linear(256, 128)),
|
||||
nn.ELU(),
|
||||
layer_init(nn.Linear(128, 1), std=1.0),
|
||||
)
|
||||
self.actor_mean = nn.Sequential(
|
||||
layer_init(nn.Linear(45 + 256, 512)),
|
||||
nn.ELU(),
|
||||
layer_init(nn.Linear(512, 256)),
|
||||
nn.ELU(),
|
||||
layer_init(nn.Linear(256, 128)),
|
||||
nn.ELU(),
|
||||
layer_init(nn.Linear(128, 12), std=0.01),
|
||||
)
|
||||
self.actor_logstd = nn.Parameter(torch.zeros(1, 12))
|
||||
|
||||
self.obs_rms = RunningMeanStd(shape = (45,))
|
||||
self.value_rms = RunningMeanStd(shape = ())
|
||||
|
||||
def forward(self, x, ac_hidden_in):
|
||||
x = self.obs_rms(x, update = False)
|
||||
ac_out, ac_hidden_out = self.actor_memory(x, ac_hidden_in)
|
||||
h = torch.cat([x, ac_out], dim = -1)
|
||||
action_mean = self.actor_mean(h)
|
||||
return action_mean, ac_hidden_out
|
||||
|
||||
class Policy:
|
||||
def __init__(self, checkpoint_path):
|
||||
self.agent = Agent()
|
||||
actor_sd = torch.load(checkpoint_path, map_location="cpu")
|
||||
self.agent.load_state_dict(actor_sd)
|
||||
onnx_file_name = checkpoint_path.replace(".pt", ".onnx")
|
||||
|
||||
dummy_input = torch.randn(1, 1, 45)
|
||||
dummy_hidden = torch.randn(1, 1, 256)
|
||||
with torch.no_grad():
|
||||
torch_out, torch_hidden_out = self.agent(dummy_input, dummy_hidden)
|
||||
|
||||
torch.onnx.export(
|
||||
self.agent, # The model being converted
|
||||
(dummy_input, dummy_hidden), # An example input for the model
|
||||
onnx_file_name, # Output file name
|
||||
export_params=True, # Store trained parameter weights inside the model file
|
||||
opset_version=11, # ONNX version (opset) to export to, adjust as needed
|
||||
do_constant_folding=True, # Whether to perform constant folding optimization
|
||||
input_names=['input', 'hidden_in'], # Name of the input in the ONNX graph (can be customized)
|
||||
output_names=['action', 'hidden_out'], # Name of the output (assuming get_action and get_value are key outputs)
|
||||
)
|
||||
|
||||
self.ort_session = onnxruntime.InferenceSession(onnx_file_name, providers=["CPUExecutionProvider"])
|
||||
ort_inputs = {'input': dummy_input.numpy(), 'hidden_in': dummy_hidden.numpy()}
|
||||
ort_outs = self.ort_session.run(None, ort_inputs)
|
||||
action, hidden_out = ort_outs[0], ort_outs[1]
|
||||
np.testing.assert_allclose(torch_out.numpy(), action, rtol=1e-03, atol=1e-05)
|
||||
np.testing.assert_allclose(torch_hidden_out.numpy(), hidden_out, rtol=1e-03, atol=1e-05)
|
||||
print("Exported model has been tested with ONNXRuntime, and the result looks good!")
|
||||
|
||||
def __call__(self, obs, info):
|
||||
#with torch.no_grad():
|
||||
# action = self.agent.get_action(self.agent.obs_rms(obs.unsqueeze(0), update = False))
|
||||
ort_inputs = {'input': obs[np.newaxis].astype(np.float32)}
|
||||
ort_outs = self.ort_session.run(None, ort_inputs)
|
||||
return ort_outs[0]
|
||||
|
||||
class CommandInterface:
|
||||
def __init__(self, limits=None):
|
||||
self.limits = limits
|
||||
self.x_vel_cmd, self.y_vel_cmd, self.yaw_vel_cmd = 0.0, 0.0, 0.0
|
||||
|
||||
def get_command(self):
|
||||
command = np.zeros((3,))
|
||||
command[0] = self.x_vel_cmd
|
||||
command[1] = self.y_vel_cmd
|
||||
command[2] = self.yaw_vel_cmd
|
||||
return command, False
|
||||
|
||||
|
||||
class CaTAgent:
|
||||
def __init__(self, command_profile, robot):
|
||||
self.robot = robot
|
||||
self.command_profile = command_profile
|
||||
# self.lcm_bridge = LCMBridgeClient(robot_name=self.robot_name)
|
||||
self.sim_dt = 0.001
|
||||
self.decimation = 20
|
||||
self.dt = self.sim_dt * self.decimation
|
||||
self.timestep = 0
|
||||
|
||||
self.device = "cpu"
|
||||
|
||||
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",
|
||||
]
|
||||
|
||||
policy_joint_names = joint_names
|
||||
|
||||
unitree_joint_names = [
|
||||
"FR_hip_joint",
|
||||
"FR_thigh_joint",
|
||||
"FR_calf_joint",
|
||||
"FL_hip_joint",
|
||||
"FL_thigh_joint",
|
||||
"FL_calf_joint",
|
||||
"RR_hip_joint",
|
||||
"RR_thigh_joint",
|
||||
"RR_calf_joint",
|
||||
"RL_hip_joint",
|
||||
"RL_thigh_joint",
|
||||
"RL_calf_joint",
|
||||
]
|
||||
|
||||
policy_to_unitree_map = []
|
||||
unitree_to_policy_map = []
|
||||
for i, policy_joint_name in enumerate(policy_joint_names):
|
||||
id = np.where([name == policy_joint_name for name in unitree_joint_names])[0][0]
|
||||
policy_to_unitree_map.append((i, id))
|
||||
self.policy_to_unitree_map = np.array(policy_to_unitree_map).astype(np.uint32)
|
||||
|
||||
for i, unitree_joint_name in enumerate(unitree_joint_names):
|
||||
id = np.where([name == unitree_joint_name for name in policy_joint_names])[0][0]
|
||||
unitree_to_policy_map.append((i, id))
|
||||
self.unitree_to_policy_map = np.array(unitree_to_policy_map).astype(np.uint32)
|
||||
|
||||
default_joint_angles = {
|
||||
"FL_hip_joint": 0.1,
|
||||
"RL_hip_joint": 0.1,
|
||||
"FR_hip_joint": -0.1,
|
||||
"RR_hip_joint": -0.1,
|
||||
"FL_thigh_joint": 0.8,
|
||||
"RL_thigh_joint": 1.0,
|
||||
"FR_thigh_joint": 0.8,
|
||||
"RR_thigh_joint": 1.0,
|
||||
"FL_calf_joint": -1.5,
|
||||
"RL_calf_joint": -1.5,
|
||||
"FR_calf_joint": -1.5,
|
||||
"RR_calf_joint": -1.5
|
||||
}
|
||||
self.default_dof_pos = np.array(
|
||||
[
|
||||
default_joint_angles[name]
|
||||
for name in joint_names
|
||||
]
|
||||
)
|
||||
|
||||
self.default_dof_pos = self.default_dof_pos
|
||||
|
||||
self.p_gains = np.zeros(12)
|
||||
self.d_gains = np.zeros(12)
|
||||
for i in range(12):
|
||||
self.p_gains[i] = 40. # 20.0
|
||||
self.d_gains[i] = 1. # 0.5
|
||||
|
||||
print(f"p_gains: {self.p_gains}")
|
||||
|
||||
self.commands = np.zeros(3)
|
||||
self.actions = np.zeros((1, 12))
|
||||
self.last_actions = np.zeros((1,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.prev_joint_acc = None
|
||||
self.torques = np.zeros(12)
|
||||
self.contact_state = np.ones(4)
|
||||
self.foot_contact_forces_mag = np.zeros(4)
|
||||
self.prev_foot_contact_forces_mag = np.zeros(4)
|
||||
self.test = 0
|
||||
|
||||
def wait_for_state(self):
|
||||
# return self.lcm_bridge.getStates(timeout=2)
|
||||
pass
|
||||
|
||||
def get_obs(self):
|
||||
cmds, reset_timer = self.command_profile.get_command()
|
||||
self.commands[:] = cmds
|
||||
|
||||
# self.state = self.wait_for_state()
|
||||
joint_state = self.robot.getJointStates()
|
||||
if joint_state is not None:
|
||||
self.gravity_vector = self.robot.getGravityInBody()
|
||||
self.prev_dof_pos = self.dof_pos.copy()
|
||||
self.dof_pos = np.array(joint_state['q'])[self.unitree_to_policy_map[:, 1]]
|
||||
self.prev_dof_vel = self.dof_vel.copy()
|
||||
self.dof_vel = np.array(joint_state['dq'])[self.unitree_to_policy_map[:, 1]]
|
||||
self.body_angular_vel = self.robot.getIMU()["gyro"]
|
||||
|
||||
try:
|
||||
self.foot_contact_forces_mag = self.robot.getFootContact()
|
||||
self.body_linear_vel = self.robot.getLinVel()
|
||||
except:
|
||||
pass
|
||||
|
||||
ob = np.concatenate(
|
||||
(
|
||||
self.body_angular_vel * 0.25,
|
||||
self.commands * np.array([2.0, 2.0, 0.25]),
|
||||
self.gravity_vector[:, 0],
|
||||
self.dof_pos * 1.0,
|
||||
#((self.dof_pos - self.prev_dof_pos) / self.dt) * 0.05,
|
||||
self.dof_vel * 0.05,
|
||||
self.last_actions[0]
|
||||
),
|
||||
axis=0,
|
||||
)
|
||||
|
||||
#return torch.tensor(ob, device=self.device).float()
|
||||
return ob
|
||||
|
||||
def publish_action(self, action, hard_reset=False):
|
||||
# command_for_robot = UnitreeLowCommand()
|
||||
#self.joint_pos_target = (
|
||||
# action[0, :12].detach().cpu().numpy() * 0.25
|
||||
#).flatten()
|
||||
self.joint_pos_target = (
|
||||
action[0, :12] * 0.25
|
||||
).flatten()
|
||||
self.joint_pos_target += self.default_dof_pos
|
||||
self.joint_vel_target = np.zeros(12)
|
||||
# command_for_robot.q_des = self.joint_pos_target
|
||||
# command_for_robot.dq_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)
|
||||
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
|
||||
# self.lcm_bridge.sendCommands(command_for_robot)
|
||||
self.robot.setCommands(self.joint_pos_target[self.policy_to_unitree_map[:, 1]],
|
||||
self.joint_vel_target[self.policy_to_unitree_map[:, 1]],
|
||||
self.p_gains[self.policy_to_unitree_map[:, 1]],
|
||||
self.d_gains[self.policy_to_unitree_map[:, 1]],
|
||||
np.zeros(12))
|
||||
|
||||
def reset(self):
|
||||
self.actions = torch.zeros((1, 12))
|
||||
self.time = time.time()
|
||||
self.timestep = 0
|
||||
return self.get_obs()
|
||||
|
||||
def step(self, actions, hard_reset=False):
|
||||
self.last_actions = self.actions[:]
|
||||
self.actions = 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()
|
||||
|
||||
joint_acc = np.abs(self.prev_dof_vel - self.dof_vel) / self.dt
|
||||
if self.prev_joint_acc is None:
|
||||
self.prev_joint_acc = np.zeros_like(joint_acc)
|
||||
joint_jerk = np.abs(self.prev_joint_acc - joint_acc) / self.dt
|
||||
self.prev_joint_acc = joint_acc.copy()
|
||||
|
||||
foot_contact_rate = np.abs(self.foot_contact_forces_mag - self.prev_foot_contact_forces_mag)
|
||||
self.prev_foot_contact_forces_mag = self.foot_contact_forces_mag.copy()
|
||||
|
||||
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, :],
|
||||
"body_linear_vel_cmd": self.commands[np.newaxis, 0:2],
|
||||
"body_angular_vel_cmd": self.commands[np.newaxis, 2:],
|
||||
"torques": self.torques,
|
||||
"foot_contact_forces_mag": self.foot_contact_forces_mag.copy(),
|
||||
"joint_acc": joint_acc[np.newaxis, :],
|
||||
"joint_jerk": joint_jerk[np.newaxis, :],
|
||||
"foot_contact_rate": foot_contact_rate[np.newaxis, :],
|
||||
}
|
||||
|
||||
self.timestep += 1
|
||||
return None, None, None, infos
|
||||
|
|
@ -158,13 +158,24 @@ class Policy:
|
|||
self.adaptation_module = torch.jit.load(
|
||||
os.path.join(checkpoint_path, "checkpoints/adaptation_module_latest.jit")
|
||||
)
|
||||
self.step_counter = 0
|
||||
self.perturbation = None
|
||||
|
||||
def __call__(self, obs, info):
|
||||
latent = self.adaptation_module.forward(obs["obs_history"].to("cpu"))
|
||||
action = self.body.forward(
|
||||
torch.cat((obs["obs_history"].to("cpu"), latent), dim=-1)
|
||||
)
|
||||
info["latent"] = latent
|
||||
with torch.no_grad():
|
||||
if self.perturbation is None:
|
||||
latent = self.adaptation_module.forward(obs["obs_history"].to("cpu"))
|
||||
self.perturbation = torch.rand_like(latent)*0.9
|
||||
else:
|
||||
latent = self.adaptation_module.forward(obs["obs_history"].to("cpu"))+self.perturbation
|
||||
if self.step_counter%50==0:
|
||||
self.perturbation = torch.rand_like(latent)*0.9
|
||||
self.step_counter +=1
|
||||
print(latent, self.perturbation)
|
||||
info["latent"] = latent
|
||||
action = self.body.forward(
|
||||
torch.cat((obs["obs_history"].to("cpu"), latent), dim=-1)
|
||||
)
|
||||
return action
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -20,7 +20,6 @@ class FrictionModel:
|
|||
|
||||
def softSign(self, u, temperature=0.1):
|
||||
return np.tanh(u/temperature)
|
||||
|
||||
|
||||
class Go2Model:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1,11 +1,27 @@
|
|||
import time
|
||||
import numpy as np
|
||||
|
||||
|
||||
class SafetyHypervisor():
|
||||
def __init__(self, robot):
|
||||
def __init__(self, robot, max_temperature=70, body_stability_idx=0.7):
|
||||
self.robot = robot
|
||||
|
||||
self.max_temperature = max_temperature
|
||||
self.normalize = lambda v: v/np.linalg.norm(v)
|
||||
self.compute_body_stability_idx = lambda v: (self.normalize(v).T@np.array([0,0,-1]))
|
||||
self.body_stability_idx = body_stability_idx
|
||||
|
||||
def unsafe(self):
|
||||
if self.robot.simulated:
|
||||
return False
|
||||
else:
|
||||
state = self.robot.getJointStates()
|
||||
temperature = state['temperature']
|
||||
if np.any(np.array(temperature)>self.max_temperature):
|
||||
print(f'One of the actuators reached max temperature of {self.max_temperature}')
|
||||
return True
|
||||
if self.compute_body_stability_idx(self.robot.getGravityInBody()) < self.body_stability_idx:
|
||||
print('Body is tilted too much! deactivating the controller')
|
||||
return True
|
||||
return False
|
||||
|
||||
def controlTimeout(self):
|
||||
|
|
|
|||
|
|
@ -425,6 +425,11 @@ class Go2Sim:
|
|||
body_v = world_R_body.T@self.data.qvel[0:3].reshape(3,1)
|
||||
return body_v
|
||||
|
||||
def getFootVelInWorld(self, site_name):
|
||||
Jp, Jr = self.getSiteJacobian(site_name)
|
||||
foot_vel = Jp@self.data.qvel.reshape(-1,1)
|
||||
return foot_vel
|
||||
|
||||
def stepHighlevel(self, vx, vy, omega_z, body_z_offset=0, step_height = 0.08, kp=[2, 0.5, 0.5], ki=[0.02, 0.01, 0.01]):
|
||||
policy_info = {}
|
||||
if self.step_counter % (self.control_dt // self.dt) == 0:
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
Loading…
Reference in New Issue