390 lines
24 KiB
Python
390 lines
24 KiB
Python
import numpy as np
|
|
from gymnasium import utils
|
|
from gymnasium.envs.mujoco import MujocoEnv
|
|
from gymnasium.spaces import Box
|
|
|
|
FRAME_SKIP = 10
|
|
|
|
|
|
class BDXEnv(MujocoEnv, utils.EzPickle):
|
|
"""
|
|
## Action space
|
|
|
|
| Num | Action | Control Min | Control Max | Name (in corresponding XML file) | Joint | Unit |
|
|
| ---- | ----------------------------------------------------------------- | ----------- | ----------- | -------------------------------- | -------- |------------- |
|
|
| 0 | Set position of right_hip_yaw | -0.58TODO | 0.58TODO | right_hip_yaw | cylinder | pos (rad) |
|
|
| 1 | Set position of right_hip_roll | -0.58TODO | 0.58TODO | right_hip_roll | cylinder | pos (rad) |
|
|
| 2 | Set position of right_hip_pitch | -0.58TODO | 0.58TODO | right_hip_pitch | cylinder | pos (rad) |
|
|
| 3 | Set position of right_knee_pitch | -0.58TODO | 0.58TODO | right_knee_pitch | cylinder | pos (rad) |
|
|
| 4 | Set position of right_ankle_pitch | -0.58TODO | 0.58TODO | right_ankle_pitch | cylinder | pos (rad) |
|
|
| 5 | Set position of left_hip_yaw | -0.58TODO | 0.58TODO | left_hip_yaw | cylinder | pos (rad) |
|
|
| 6 | Set position of left_hip_roll | -0.58TODO | 0.58TODO | left_hip_roll | cylinder | pos (rad) |
|
|
| 7 | Set position of left_hip_pitch | -0.58TODO | 0.58TODO | left_hip_pitch | cylinder | pos (rad) |
|
|
| 8 | Set position of left_knee_pitch | -0.58TODO | 0.58TODO | left_knee_pitch | cylinder | pos (rad) |
|
|
| 9 | Set position of left_ankle_pitch | -0.58TODO | 0.58TODO | left_ankle_pitch | cylinder | pos (rad) |
|
|
| 9 | Set position of head_pitch1 | -0.58TODO | 0.58TODO | left_ankle_pitch | cylinder | pos (rad) |
|
|
| 9 | Set position of head_pitch2 | -0.58TODO | 0.58TODO | left_ankle_pitch | cylinder | pos (rad) |
|
|
| 9 | Set position of head_yaw | -0.58TODO | 0.58TODO | left_ankle_pitch | cylinder | pos (rad) |
|
|
|
|
|
|
## Observation space
|
|
|
|
| Num | Observation | Min | Max | Name (in corresponding XML file) | Joint | Unit |
|
|
| --- | -------------------------------------------------------- | ---- | --- | -------------------------------- | -------- | ------------------------ |
|
|
| 0 | Rotation right_hip_yaw | -Inf | Inf | right_hip_yaw | cylinder | angle (rad) |
|
|
| 1 | Rotation right_hip_roll | -Inf | Inf | right_hip_roll | cylinder | angle (rad) |
|
|
| 2 | Rotation right_hip_pitch | -Inf | Inf | right_hip_pitch | cylinder | angle (rad) |
|
|
| 3 | Rotation right_knee_pitch | -Inf | Inf | right_knee_pitch | cylinder | angle (rad) |
|
|
| 4 | Rotation right_ankle_pitch | -Inf | Inf | right_ankle_pitch | cylinder | angle (rad) |
|
|
| 5 | Rotation left_hip_yaw | -Inf | Inf | left_hip_yaw | cylinder | angle (rad) |
|
|
| 6 | Rotation left_hip_roll | -Inf | Inf | left_hip_roll | cylinder | angle (rad) |
|
|
| 7 | Rotation left_hip_pitch | -Inf | Inf | left_hip_pitch | cylinder | angle (rad) |
|
|
| 8 | Rotation left_knee_pitch | -Inf | Inf | left_knee_pitch | cylinder | angle (rad) |
|
|
| 9 | Rotation left_ankle_pitch | -Inf | Inf | left_ankle_pitch | cylinder | angle (rad) |
|
|
| 10 | Rotation neck_pitch | -Inf | Inf | head_pitch1 | cylinder | angle (rad) |
|
|
| 11 | Rotation head_pitch | -Inf | Inf | head_pitch2 | cylinder | angle (rad) |
|
|
| 12 | Rotation head_yaw | -Inf | Inf | head_yaw | cylinder | angle (rad) |
|
|
| 13 | Velocity of right_hip_yaw | -Inf | Inf | right_hip_yaw | cylinder | speed (rad/s) |
|
|
| 14 | Velocity of right_hip_roll | -Inf | Inf | right_hip_roll | cylinder | speed (rad/s) |
|
|
| 15 | Velocity of right_hip_pitch | -Inf | Inf | right_hip_pitch | cylinder | speed (rad/s) |
|
|
| 16 | Velocity of right_knee_pitch | -Inf | Inf | right_knee_pitch | cylinder | speed (rad/s) |
|
|
| 17 | Velocity of right_ankle_pitch | -Inf | Inf | right_ankle_pitch | cylinder | speed (rad/s) |
|
|
| 18 | Velocity of left_hip_yaw | -Inf | Inf | left_hip_yaw | cylinder | speed (rad/s) |
|
|
| 19 | Velocity of left_hip_roll | -Inf | Inf | left_hip_roll | cylinder | speed (rad/s) |
|
|
| 20 | Velocity of left_hip_pitch | -Inf | Inf | left_hip_pitch | cylinder | speed (rad/s) |
|
|
| 21 | Velocity of left_knee_pitch | -Inf | Inf | left_knee_pitch | cylinder | speed (rad/s) |
|
|
| 22 | Velocity of left_ankle_pitch | -Inf | Inf | left_ankle_pitch | cylinder | speed (rad/s) |
|
|
| 23 | Velocity of neck_pitch | -Inf | Inf | head_pitch1 | cylinder | speed (rad/s) |
|
|
| 24 | Velocity of head_pitch | -Inf | Inf | head_pitch2 | cylinder | speed (rad/s) |
|
|
| 25 | Velocity of head_yaw | -Inf | Inf | head_yaw | cylinder | speed (rad/s) |
|
|
|
|
|
|
| 26 | roll angular velocity | -Inf | Inf | | | |e
|
|
| 27 | pitch angular velocity | -Inf | Inf | | | |
|
|
| 28 | yaw angular velocity | -Inf | Inf | | | |
|
|
| 29 | current x linear velocity | -Inf | Inf | | | |
|
|
| 30 | current y linear velocity | -Inf | Inf | | | |
|
|
| 31 | current z linear velocity | -Inf | Inf | | | |
|
|
| 32 | current x target linear velocity | -Inf | Inf | | | |
|
|
| 33 | current y target linear velocity | -Inf | Inf | | | |
|
|
| 34 | current yaw target angular velocity | -Inf | Inf | | | |
|
|
|
|
Changed to control history
|
|
| 35 | t-1 right_hip_yaw rotation error | -Inf | Inf | | | |
|
|
| 36 | t-1 right_hip_roll rotation error | -Inf | Inf | | | |
|
|
| 37 | t-1 right_hip_pitch rotation error | -Inf | Inf | | | |
|
|
| 38 | t-1 right_knee_pitch rotation error | -Inf | Inf | | | |
|
|
| 39 | t-1 right_ankle_pitch rotation error | -Inf | Inf | | | |
|
|
| 40 | t-1 left_hip_yaw rotation error | -Inf | Inf | | | |
|
|
| 41 | t-1 left_hip_roll rotation error | -Inf | Inf | | | |
|
|
| 42 | t-1 left_hip_pitch rotation error | -Inf | Inf | | | |
|
|
| 43 | t-1 left_knee_pitch rotation error | -Inf | Inf | | | |
|
|
| 44 | t-1 left_ankle_pitch rotation error | -Inf | Inf | | | |
|
|
| 45 | t-1 neck_pitch rotation error | -Inf | Inf | | | |
|
|
| 46 | t-1 head_pitch rotation error | -Inf | Inf | | | |
|
|
| 47 | t-1 head_yaw rotation error | -Inf | Inf | | | |
|
|
|
|
| 48 | t-2 right_hip_yaw rotation error | -Inf | Inf | | | |
|
|
| 49 | t-2 right_hip_roll rotation error | -Inf | Inf | | | |
|
|
| 50 | t-2 right_hip_pitch rotation error | -Inf | Inf | | | |
|
|
| 51 | t-2 right_knee_pitch rotation error | -Inf | Inf | | | |
|
|
| 52 | t-2 right_ankle_pitch rotation error | -Inf | Inf | | | |
|
|
| 53 | t-2 left_hip_yaw rotation error | -Inf | Inf | | | |
|
|
| 54 | t-2 left_hip_roll rotation error | -Inf | Inf | | | |
|
|
| 55 | t-2 left_hip_pitch rotation error | -Inf | Inf | | | |
|
|
| 56 | t-2 left_knee_pitch rotation error | -Inf | Inf | | | |
|
|
| 57 | t-2 left_ankle_pitch rotation error | -Inf | Inf | | | |
|
|
| 58 | t-2 neck_pitch rotation error | -Inf | Inf | | | |
|
|
| 59 | t-2 head_pitch rotation error | -Inf | Inf | | | |
|
|
| 60 | t-2 head_yaw rotation error | -Inf | Inf | | | |
|
|
| 61 | t-3 right_hip_yaw rotation error | -Inf | Inf | | | |
|
|
| 62 | t-3 right_hip_roll rotation error | -Inf | Inf | | | |
|
|
| 63 | t-3 right_hip_pitch rotation error | -Inf | Inf | | | |
|
|
| 64 | t-3 right_knee_pitch rotation error | -Inf | Inf | | | |
|
|
| 65 | t-3 right_ankle_pitch rotation error | -Inf | Inf | | | |
|
|
| 66 | t-3 left_hip_yaw rotation error | -Inf | Inf | | | |
|
|
| 67 | t-3 left_hip_roll rotation error | -Inf | Inf | | | |
|
|
| 68 | t-3 left_hip_pitch rotation error | -Inf | Inf | | | |
|
|
| 69 | t-3 left_knee_pitch rotation error | -Inf | Inf | | | |
|
|
| 70 | t-3 left_ankle_pitch rotation error | -Inf | Inf | | | |
|
|
| 71 | t-3 neck_pitch rotation error | -Inf | Inf | | | |
|
|
| 72 | t-3 head_pitch rotation error | -Inf | Inf | | | |
|
|
| 73 | t-3 head_yaw rotation error | -Inf | Inf | | | |
|
|
|
|
| 74 | left foot in contact with the floor | -Inf | Inf | | | |
|
|
| 75 | right foot in contact with the floor | -Inf | Inf | | | |
|
|
|
|
| x76 | t | -Inf | Inf | | | |
|
|
| x74 | sinus | -Inf | Inf | | | |
|
|
|
|
"""
|
|
|
|
metadata = {
|
|
"render_modes": [
|
|
"human",
|
|
"rgb_array",
|
|
"depth_array",
|
|
],
|
|
"render_fps": 50,
|
|
}
|
|
|
|
# Ideas
|
|
# Low pass filter on the joint angles
|
|
|
|
def __init__(self, **kwargs):
|
|
utils.EzPickle.__init__(self, **kwargs)
|
|
observation_space = Box(low=-np.inf, high=np.inf, shape=(86,), dtype=np.float64)
|
|
self.target_velocity = np.asarray([0, 0, 0]) # x, y, yaw
|
|
self.joint_history_length = 3
|
|
self.joint_error_history = self.joint_history_length * [15 * [0]]
|
|
self.joint_ctrl_history = self.joint_history_length * [15 * [0]]
|
|
|
|
self.left_foot_in_contact = 0
|
|
self.right_foot_in_contact = 0
|
|
self.startup_cooldown = 1.0
|
|
|
|
self.prev_t = 0
|
|
self.init_pos = np.array(
|
|
[
|
|
-0.013946457213457239,
|
|
0.07918837709879874,
|
|
0.5325073962634973,
|
|
-1.6225192902713386,
|
|
0.9149246381274986,
|
|
0.013627156377842975,
|
|
0.07738878096596595,
|
|
0.5933527914082196,
|
|
-1.630548419252953,
|
|
0.8621333440557593,
|
|
-0.17453292519943295,
|
|
-0.17453292519943295,
|
|
8.65556854322817e-27,
|
|
0,
|
|
0,
|
|
]
|
|
)
|
|
MujocoEnv.__init__(
|
|
self,
|
|
"/home/antoine/MISC/mini_BDX/mini_bdx/robots/bdx/scene.xml",
|
|
FRAME_SKIP,
|
|
observation_space=observation_space,
|
|
**kwargs,
|
|
)
|
|
|
|
def check_contact(self, body1_name, body2_name):
|
|
body1_id = self.data.body(body1_name).id
|
|
body2_id = self.data.body(body2_name).id
|
|
|
|
for i in range(self.data.ncon):
|
|
contact = self.data.contact[i]
|
|
|
|
if (
|
|
self.model.geom_bodyid[contact.geom1] == body1_id
|
|
and self.model.geom_bodyid[contact.geom2] == body2_id
|
|
) or (
|
|
self.model.geom_bodyid[contact.geom1] == body2_id
|
|
and self.model.geom_bodyid[contact.geom2] == body1_id
|
|
):
|
|
return True
|
|
|
|
return False
|
|
|
|
def feet_contact_reward(self):
|
|
return self.right_foot_in_contact + self.left_foot_in_contact
|
|
|
|
def velocity_tracking_reward(self):
|
|
base_velocity = list(self.data.body("base").cvel[3:][:2]) + [
|
|
self.data.body("base").cvel[:3][2]
|
|
]
|
|
base_velocity = np.asarray(base_velocity)
|
|
return np.exp(-np.square(base_velocity - self.target_velocity).sum())
|
|
|
|
def smoothness_reward(self):
|
|
# Warning, this function only works if the history is 3 :)
|
|
smooth = 0
|
|
t0 = self.joint_ctrl_history[0]
|
|
t_minus1 = self.joint_ctrl_history[1]
|
|
t_minus2 = self.joint_ctrl_history[2]
|
|
|
|
for i in range(15):
|
|
smooth += 2.5 * np.square(t0[i] - t_minus1[i]) + 1.5 * np.square(
|
|
t0[i] - 2 * t_minus1[i] + t_minus2[i]
|
|
)
|
|
|
|
return -smooth
|
|
|
|
def smoothness_reward2(self):
|
|
# Warning, this function only works if the history is 3 :)
|
|
smooth = 0
|
|
t0 = self.joint_ctrl_history[0]
|
|
t_minus1 = self.joint_ctrl_history[1]
|
|
t_minus2 = self.joint_ctrl_history[2]
|
|
|
|
for i in range(15):
|
|
smooth += np.square(t0[i] - t_minus1[i]) + np.square(
|
|
t_minus1[i] - t_minus2[i]
|
|
)
|
|
# smooth += 2.5 * np.square(t0[i] - t_minus1[i]) + 1.5 * np.square(
|
|
# t0[i] - 2 * t_minus1[i] + t_minus2[i]
|
|
# )
|
|
|
|
return -smooth
|
|
|
|
def joint_velocity_reward(self):
|
|
return -np.square(self.data.qvel[:]).sum()
|
|
|
|
def walking_height_reward(self):
|
|
return (
|
|
-np.square((self.get_body_com("base")[2] - 0.15)) * 100
|
|
) # "normal" walking height is about 0.15m
|
|
|
|
def upright_reward(self):
|
|
# angular distance to upright position in reward
|
|
Z_vec = np.array(self.data.body("base").xmat).reshape(3, 3)[:, 2]
|
|
return np.square(np.dot(np.array([0, 0, 1]), Z_vec))
|
|
|
|
def init_position_reward(self):
|
|
return -np.square(self.data.qpos[7 : 7 + 15] - self.init_pos).sum()
|
|
|
|
def is_terminated(self) -> bool:
|
|
rot = np.array(self.data.body("base").xmat).reshape(3, 3)
|
|
Z_vec = rot[:, 2]
|
|
upright = np.array([0, 0, 1])
|
|
return (
|
|
self.data.body("base").xpos[2] < 0.08 or np.dot(upright, Z_vec) <= 0
|
|
) # base z is below 0.08m or base has more than 90 degrees of tilt
|
|
|
|
def action_LFP(self, action):
|
|
# Low pass filter for the actions
|
|
action_tminus1 = self.joint_ctrl_history[0]
|
|
action_tminus2 = self.joint_ctrl_history[1]
|
|
action_tminus3 = self.joint_ctrl_history[2]
|
|
# return action * 0.5 + action_tminus1 * 0.5
|
|
return (
|
|
action * 0.5
|
|
+ action_tminus1 * 0.3
|
|
+ action_tminus2 * 0.15
|
|
+ action_tminus3 * 0.05
|
|
)
|
|
|
|
d_action = action - action_tminus1
|
|
|
|
action = [
|
|
a if abs(d) < 0.1 else b
|
|
for a, b, d in zip(action, action_tminus1, d_action)
|
|
]
|
|
return action
|
|
|
|
def step(self, a):
|
|
# https://www.nature.com/articles/s41598-023-38259-7.pdf
|
|
|
|
# add random force
|
|
# if np.random.rand() < 0.05:
|
|
# self.data.xfrc_applied[self.data.body("base").id][:3] = [
|
|
# np.random.randint(-5, 5),
|
|
# np.random.randint(-5, 5),
|
|
# np.random.randint(-5, 5),
|
|
# ] # absolute
|
|
|
|
t = self.data.time
|
|
dt = t - self.prev_t
|
|
if self.startup_cooldown > 0:
|
|
self.startup_cooldown -= dt
|
|
|
|
if self.startup_cooldown > 0:
|
|
self.do_simulation(self.init_pos, FRAME_SKIP)
|
|
reward = 0
|
|
else:
|
|
self.do_simulation(a, FRAME_SKIP)
|
|
# self.do_simulation(self.action_LFP(a), 4)
|
|
self.right_foot_in_contact = self.check_contact("foot_module", "floor")
|
|
self.left_foot_in_contact = self.check_contact("foot_module_2", "floor")
|
|
|
|
reward = (
|
|
0.05 # time reward
|
|
# + 0.1 * self.walking_height_reward()
|
|
# + 0.1 * self.upright_reward()
|
|
# + 0.1 * self.velocity_tracking_reward()
|
|
+ 0.01 * self.smoothness_reward2()
|
|
+ 0.1 * self.init_position_reward()
|
|
# + 0.1 * self.feet_contact_reward()
|
|
# + 0.1 * self.joint_angle_deviation_reward()
|
|
# + 0.01 * self.follow_walk_engine_reward(dt)
|
|
# + 0.001 * self.joint_velocity_reward()
|
|
)
|
|
|
|
ob = self._get_obs()
|
|
|
|
if self.render_mode == "human":
|
|
self.render()
|
|
|
|
self.prev_t = t
|
|
|
|
return (ob, reward, self.is_terminated(), False, {}) # terminated # truncated
|
|
|
|
def reset_model(self):
|
|
self.prev_t = self.data.time
|
|
self.startup_cooldown = 1.0
|
|
|
|
# self.model.opt.gravity[:] = [0, 0, 0] # no gravity
|
|
|
|
self.goto_init()
|
|
|
|
self.joint_error_history = self.joint_history_length * [15 * [0]]
|
|
self.joint_ctrl_history = self.joint_history_length * [15 * [0]]
|
|
# qpos = self.data.qpos
|
|
|
|
# LATEST
|
|
# added randomization to the initial position
|
|
# for i in range(7, len(qpos)):
|
|
# qpos[i] = np.random.uniform(-0.3, 0.3)
|
|
|
|
# self.init_qpos = qpos.copy().flatten()
|
|
|
|
# Randomize later
|
|
# self.target_velocity = np.asarray([0, 0, 0]) # x, y, yaw
|
|
self.target_velocity = np.asarray([0, 0, 0]) # x, y, yaw
|
|
|
|
self.set_state(self.data.qpos, self.data.qvel)
|
|
return self._get_obs()
|
|
|
|
def goto_init(self):
|
|
self.data.qpos[:] = np.zeros(len(self.data.qpos[:]))
|
|
self.data.qpos[2] = 0.15
|
|
# self.data.qpos[3 : 3 + 4] = [1, 0, 0, 0]
|
|
self.data.qpos[3 : 3 + 4] = [1, 0, 0.08, 0]
|
|
|
|
self.data.qvel[:] = np.zeros(len(self.data.qvel[:]))
|
|
self.data.qpos[7 : 7 + 15] = self.init_pos
|
|
self.data.ctrl[:] = self.init_pos
|
|
|
|
def _get_obs(self):
|
|
joints_rotations = self.data.qpos[7 : 7 + 15]
|
|
joints_velocities = self.data.qvel[6 : 6 + 15]
|
|
|
|
# joints_error = self.data.ctrl - self.data.qpos[7 : 7 + 13]
|
|
# self.joint_error_history.append(joints_error)
|
|
# self.joint_error_history = self.joint_error_history[
|
|
# -self.joint_history_length :
|
|
# ]
|
|
|
|
angular_velocity = self.data.body("base").cvel[
|
|
:3
|
|
] # TODO this is imu, add noise to it later
|
|
linear_velocity = self.data.body("base").cvel[3:]
|
|
|
|
self.joint_ctrl_history.append(self.data.ctrl.copy())
|
|
self.joint_ctrl_history = self.joint_ctrl_history[-self.joint_history_length :]
|
|
|
|
return np.concatenate(
|
|
[
|
|
joints_rotations,
|
|
joints_velocities,
|
|
angular_velocity,
|
|
linear_velocity,
|
|
self.target_velocity,
|
|
np.array(self.joint_ctrl_history).flatten(),
|
|
[self.left_foot_in_contact, self.right_foot_in_contact],
|
|
]
|
|
)
|