JiajunLI/roboimi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore(project): 项目初始化

Browse Source
This commit is contained in:
JiajunLI
2026-01-27 15:42:02 +08:00
commit 37cc693abb
359 changed files with 264112 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

242
roboimi/utils/KDL_utils/lma_kdl.py Normal file
View File

@@ -0,0 +1,242 @@
import numpy as np
import time
from roboimi.utils.KDL_utils import KDL_utils
from roboimi.utils.gym_wrapper import GymWrapper
from stable_baselines3 import PPO, SAC, TD3
from stable_baselines3.common.callbacks import BaseCallback
from gym.wrappers import TimeLimit
# URDF_PATH = "/home/zhouhr/mujoco_diana/roboIMI/roboimi/assets/models/manipulators/DianaMed/DianaMed.urdf"
URDF_PATH = "/home/leeezd/tmp/RoboIMI/roboimi/assets/models/manipulators/DianaMed/DianaMed.urdf"
class Lma_hcm(object):
def __init__(self, verbose=False):
self.kin = KDL_utils(URDF_PATH)
self.verbose = verbose
# robot parameters
self.Dof = 7
self.init_qpos = np.zeros(self.Dof)
self.current_qpos = np.zeros(self.Dof)
self.target_qpos = np.zeros(self.Dof)
# hyper-params
self.max_episode_steps = 100
self.eps_phi = 1e-6
self.lamb = 0.01 # init
self.name = "lma_kdl"
self.v_d = 1.2 # default
self.v_m = 1.2
# lma params
self.iter_times = 0
self.prev_phi = None
# training params
self.action_dim = (2,)
self.obs_dim = (2,)
self.action_low = np.array([-1.0, 1.0])
self.action_high = np.array([1.0, 1.0])
def setInitQpos(self, qpos: np.array):
self.init_qpos = qpos
self.current_qpos = qpos
def setTargetQpos(self, qpos: np.array):
self.target_qpos = qpos
def computeError(self, current_joints):
""" Phi means the error between current pose and target pose, eg. e = Pd - P """
current_T = self.kin.getEEtf(self.kin.setJntArray(current_joints))
current_pos = np.array([current_T.p[0], current_T.p[1], current_T.p[2]])
current_rot = current_T.M
target_T = self.kin.getEEtf(self.kin.setJntArray(self.target_qpos))
target_pos = np.array([target_T.p[0], target_T.p[1], target_T.p[2]])
target_rot = target_T.M
current_pos_err = target_pos - current_pos
current_rot_err = (target_rot * (current_rot.Inverse())).GetRot()
err = np.array([current_pos_err[0], current_pos_err[1], current_pos_err[2],
current_rot_err[0], current_rot_err[1], current_rot_err[2]])
return err
def computePhi(self, err):
return err.T @ np.eye(6) @ err
def computeDelta(self, lamb, qpos, err):
jac = self.kin.setNumpyMat(self.kin.getJac(self.kin.setJntArray(qpos)))
mat = jac.T.dot(jac) + lamb * np.eye(7)
if np.linalg.det(mat) == 0:
raise np.linalg.LinAlgError
inv_mat = np.linalg.inv(mat)
return inv_mat @ jac.T @ err
def updatePhi(self, lamb):
current_qpos = self.current_qpos.copy()
err = self.computeError(current_qpos)
delta_q = self.computeDelta(lamb, current_qpos, err)
phi = self.computePhi(err)
new_qpos = current_qpos + delta_q
return phi, lamb, new_qpos
def updateLambda(self):
# phi, lamb, current_qpos = self.update(self.lamb * np.exp(-self.v))
phi, lamb, qpos = self.updatePhi(self.lamb / (self.v_d + 1e-5))
if phi - self.prev_phi < self.eps_phi:
return lamb, qpos, phi, True
phi, lamb, qpos = self.updatePhi(self.lamb)
if phi - self.prev_phi < self.eps_phi:
return lamb, qpos, phi, True
# phi, lamb, qpos = self.update(self.lamb * np.exp(self.v))
phi, lamb, qpos = self.updatePhi(self.lamb * self.v_m)
if phi - self.prev_phi < self.eps_phi:
return lamb, qpos, phi, True
# success = False
# while phi - self.prev_phi >= 0:
# _phi, _lamb, _qpos = self.updatePhi(lamb * self.v_m)
# if phi - self.prev_phi > 0:
# print("lamb:", _lamb)
# break
# phi = _phi
# lamb = _lamb
# qpos = _qpos
# Failure
return self.lamb, self.current_qpos, self.prev_phi, False
def step(self, action):
""" LMA computation """
self.iter_times += 1
self.v_d = (1 / 1 + np.exp(action[0])) * 100
self.v_m = (1 / 1 + np.exp(action[1])) * 100
done = False
reward = 0.0
step_start_time = time.time()
self.lamb, self.current_qpos, self.prev_phi, success = self.updateLambda()
step_cost_time = time.time() - step_start_time
print(
f"times={self.iter_times}, phi={self.prev_phi:.7f}, cost_time={step_cost_time:.5f}. lamb:{self.lamb}")
# reward += 0.3 * (1 - 1e3 * step_cost_time) if (1 - 1e3 * step_cost_time) >= 0 else 0.0
reward -= 0.3 * self.iter_times
if self.prev_phi >= 0.1:
reward += 0.5 * (1.0 - 1 * self.prev_phi)
elif 0.1 > self.prev_phi >= 0.01:
reward += 0.5 + (1.0 - 10 * self.prev_phi)
elif 0.01 > self.prev_phi >= 0.001:
reward += 1.5 + (1.0 - 100 * self.prev_phi)
elif 0.001 > self.prev_phi >= 0.0001:
reward += 2.5 + (1.0 - 1000 * self.prev_phi)
elif 0.0001 > self.prev_phi >= 0.00001:
reward += 3.5 + (1.0 - 10000 * self.prev_phi)
elif 0.00001 > self.prev_phi >= 0.000001:
reward += 4.5 + (1.0 - 10000 * self.prev_phi)
else:
# if self.prev_phi <= self.eps_phi:
done = True
reward += 10.0
if success is False:
reward = 0
return self.get_obs(), reward, done, dict()
def reset(self):
self.setInitQpos(3.14 * np.random.uniform(low=-1.0, high=1.0, size=(7,)))
self.setTargetQpos(3.14 * np.random.uniform(low=-1.0, high=1.0, size=(7,)))
err = self.computeError(self.init_qpos)
self.prev_phi = self.computePhi(err)
self.iter_times = 0
self.lamb = 0.01
print("=====================================new episode==================================")
if self.verbose:
print("init phi =", self.prev_phi)
print("init qpos =", self.init_qpos)
print("goal qpos =", self.target_qpos)
result = self.kin.getEEtf(self.kin.setJntArray(self.current_qpos))
target = self.kin.getEEtf(self.kin.setJntArray(self.target_qpos))
print("init T =", result)
print("goal T =", target)
return self.get_obs()
def get_obs(self):
obs = np.zeros(self.obs_dim)
obs[0] = np.array([self.prev_phi], dtype=np.float32)
obs[1] = np.array([self.lamb], dtype=np.float32)
# obs[1:8] = self.current_qpos
# obs[8:15] = self.target_qpos
return obs
def run_test():
env = Lma_hcm(verbose=False)
env.reset()
for n_steps in range(int(1e6)):
a = np.array([0.5, 0.1])
s_, r, d, _ = env.step(a)
# print(r, d)
if d or env.iter_times % env.max_ep_steps == 0:
if env.verbose:
result = env.kin.getEEtf(env.kin.setJntArray(env.current_qpos))
target = env.kin.getEEtf(env.kin.setJntArray(env.target_qpos))
print("compute_results =", result)
print("target_T = ", target)
env.reset()
def run_train():
env = Lma_hcm()
env = GymWrapper(env)
env = TimeLimit(env, max_episode_steps=env.env.max_episode_steps)
# check_env(env)
class TensorboardCallback(BaseCallback):
"""
Custom callback for plotting additional values in tensorboard.
"""
def __init__(self, log_dir, verbose=0):
super(TensorboardCallback, self).__init__(verbose)
self.log_dir = log_dir
def _on_step(self) -> bool:
if self.n_calls % 51200 == 0:
print("Saving new best model")
self.model.save(self.log_dir + f"/model_saved/PPO/LMA_{self.n_calls}")
return True
log_dir = "log/"
model = PPO(
policy="MlpPolicy",
env=env,
tensorboard_log=log_dir,
device="cuda",
)
model.learn(total_timesteps=int(1e7), callback=TensorboardCallback(log_dir=log_dir))
model.save("lma_rl")
def run_play():
env = Lma_hcm(verbose=False)
model = PPO.load("log/model_saved/PPO_2/LMA_3788800.zip")
obs = env.reset()
for i in range(int(1e6)):
action, _states = model.predict(obs)
obs, reward, done, info = env.step(action)
print("a", action)
if done:
env.reset()
if __name__ == '__main__':
# run_test()
run_train()
# run_play()