meiqi/tienkung-szu
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

First Commit

Browse Source
This commit is contained in:
meiqi
2026-03-27 16:10:51 +08:00
commit c45245038f
103 changed files with 10994 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

98
Deploy_Tienkung/policy/walk_amp/config/walk_amp.yaml Normal file
View File

@@ -0,0 +1,98 @@
model_path: "policy.onnx"
motor_num: 29 # 电机数量
actions_size: 23 # action的大小
dt: 0.01
warm_start_time: 0.0
xsense_data_roll_offset: 0.0
joint_names: [
hip_pitch_l_joint, hip_pitch_r_joint, waist_yaw_joint,
hip_roll_l_joint, hip_roll_r_joint, waist_roll_joint,
hip_yaw_l_joint, hip_yaw_r_joint, waist_pitch_joint,
knee_pitch_l_joint, knee_pitch_r_joint,
shoulder_pitch_l_joint, shoulder_pitch_r_joint,
ankle_pitch_l_joint, ankle_pitch_r_joint,
shoulder_roll_l_joint, shoulder_roll_r_joint,
ankle_roll_l_joint, ankle_roll_r_joint,
shoulder_yaw_l_joint, shoulder_yaw_r_joint,
elbow_pitch_l_joint, elbow_pitch_r_joint
]
zero_pos_offset: [
0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 0.0, 0.0,
0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
]
ct_scale: [
2.1, 2.1, 3.207, 2.673, 2.6, 2.6,
2.1, 2.1, 3.207, 2.673, 2.6, 2.6,
3.207, 3.207, 3.207, #TODO:待更新 pry
3.207, 2.28, 5.89, 5.89, 3.35, 2.3, 2.3,
3.207, 2.28, 5.89, 5.89, 3.35, 2.3, 2.3
]
control:
# gait_cycle_period: 0.9
action_scale: 0.25 # 动作缩放比例，可写标量或 23 维数组
decimation: 1 # 策略下发频率控制
gait:
gait_air_ratio_l: 0.38
gait_air_ratio_r: 0.38
gait_phase_offset_l: 0.38
gait_phase_offset_r: 0.88
gait_cycle: 0.85
normalization:
clip_scales:
clip_observations: 100. #18.0mlp
clip_actions: 100. #18.0
obs_scales:
lin_vel: 1.0 # 线速度缩放因子
ang_vel: 1.0 # 角速度缩放因子
dof_pos: 1.0 # 关节位置缩放因子
dof_vel: 1.0 # 关节速度缩放因子
size:
num_hist: 10 # 历史观测帧数
observations_size: 84 # 单帧观测长度
gains:
kp: [
300.0, 300.0, 400.0,
300.0, 300.0, 400.0,
150.0, 150.0, 400.0,
350.0, 350.0,
150.0, 150.0,
30.0, 30.0,
50.0, 50.0,
16.8, 16.8,
50.0, 50.0,
150.0, 150.0
]
kd: [
10.0, 10.0, 5.0,
10.0, 10.0, 10.0,
5.0, 5.0, 10.0,
10.0, 10.0,
7.5, 7.5,
2.5, 2.5,
2.5, 2.5,
1.4, 1.4,
2.5, 2.5,
5.0, 5.0
]
init_state:
default_joint_angles: [
-0.15, -0.15, 0.0,
0.0, 0.0, 0.0,
-0.0, -0.0, 0.0,
0.3, 0.3,
0.2, 0.2,
-0.15, -0.15,
0.1, -0.1,
0.0, 0.0,
0.0, 0.0,
-0.5, -0.5
]

454
Deploy_Tienkung/policy/walk_amp/fsm_walkamp.py Normal file
View File

@@ -0,0 +1,454 @@
"""
FSM State Implementations
Concrete implementations of different FSM states
"""
import numpy as np
import onnxruntime as ort
from FSM.fsm_base import FSMState, FSMStateName
from common.joystick import ControlFlag
from common.robot_data import RobotData
from common.BasicFunction import clip_vector, gait_phase
import os
import yaml
from scipy.spatial.transform import Rotation
class FSMStateWALKAMP(FSMState):
"""WALKAMP策略状态实现"""
def _reset_internal_state(self):
"""把所有随时间变化的内部状态重置成初始值"""
# 1) 清空 obs / hist / actions
self.observations_.fill(0.0)
self.proprio_hist_buf_.fill(0.0)
self.last_actions_.fill(0.0)
self.actions_.fill(0.0)
# 2) 标志位重置
self.is_first_obs_ = True
self.is_first_action_ = True
self.is_first_step_ = True
# 3) 期望关节 / 期望速度 / 力矩重置为“初始姿态”
# 你已经有 self.joint_pos_array（mj 顺序，长度 len(self.joint_xml)）
base = self.robot_data_.q_d_.shape[0] - self.motor_num_
# 期望角 = 初始角
self.robot_data_.q_d_[base:base + len(self.joint_xml)] = self.joint_pos_array
# 期望速度 = 0
self.robot_data_.q_dot_d_[base:base + len(self.joint_xml)] = 0.0
# 期望力矩 = 0（位置控制）
self.robot_data_.tau_d_[base:base + len(self.joint_xml)] = 0.0
def __init__(self,
robot_data: RobotData,
config_path: str | None = None,
base_dir: str | None = None,
log_prefix: str = "FSMStateWALKAMP"):
super().__init__(robot_data)
self.log_prefix = log_prefix
# 获取包路径
current_dir = os.path.dirname(os.path.abspath(__file__))
if config_path is None:
config_path = os.path.join(current_dir, "config", "walk_amp.yaml")
config_path = os.path.abspath(config_path)
if base_dir is None:
base_dir = os.path.dirname(os.path.dirname(config_path))
with open(config_path, 'r') as f:
policy_config = yaml.safe_load(f)
# Load configuration exactly like C++
self.action_num_ = policy_config.get('actions_size')
self.motor_num_ = policy_config.get('motor_num')
self.dt_ = policy_config.get('dt')
# Size configuration
size_config = policy_config.get('size', {})
self.num_hist_ = size_config.get('num_hist')
self.obs_size_ = size_config.get('observations_size')
# Control configuration
control_config = policy_config.get('control', {})
self.action_scale_ = control_config.get('action_scale')
# self.gait_cycle_period_ = control_config.get('gait_cycle_period', 1.0)
self.decimation_ = control_config.get('decimation')
self.warm_start_time_ = control_config.get(
'warm_start_time',
policy_config.get('warm_start_time', 0.3),
)
# Normalization configuration
norm_config = policy_config.get('normalization', {})
clip_config = norm_config.get('clip_scales', {})
obs_config = norm_config.get('obs_scales', {})
self.clip_obs_ = clip_config.get('clip_observations', 100.0)
self.clip_act_ = clip_config.get('clip_actions', 100.0)
self.lin_vel_scale_ = obs_config.get('lin_vel')
self.ang_vel_scale_ = obs_config.get('ang_vel')
self.dof_pos_scale_ = obs_config.get('dof_pos')
self.dof_vel_scale_ = obs_config.get('dof_vel')
# Initialize buffers and actions
self.observations_ = np.zeros(self.obs_size_ * self.num_hist_, dtype=np.float32)
self.proprio_hist_buf_ = np.zeros(self.obs_size_ * self.num_hist_, dtype=np.float32)
self.last_actions_ = np.zeros(self.action_num_, dtype=np.float32)
self.actions_ = np.zeros(self.action_num_, dtype=np.float32)
self._warm_start_pose = np.zeros(self.motor_num_, dtype=np.float32)
# Flags matching C++
self.is_first_obs_ = True
self.is_first_action_ = True
# self.phase_locked = False
self.timer_gait_ = 0.0
gait_config = policy_config.get('gait', {})
self.gait_cycle = gait_config.get('gait_cycle', 0.85)
self.left_phase_ratio = gait_config.get('gait_air_ratio_l', 0.38)
self.right_phase_ratio = gait_config.get('gait_air_ratio_r', 0.38)
self.left_theta_offset = gait_config.get('gait_phase_offset_l', 0.38)
self.right_theta_offset = gait_config.get('gait_phase_offset_r', 0.88)
self.is_first_step_ = True
step = (self.decimation_ if self.decimation_ else 1) * self.dt_
if self.warm_start_time_ > 0 and step > 0:
self._warm_start_steps = max(1, int(self.warm_start_time_ / step))
else:
self._warm_start_steps = 0
self._warmup_inference_counter = 0
# Initialize ONNX session
self.model_path = os.path.join(base_dir, "model", policy_config["model_path"])
self._init_onnx_session()
self.joint_seq = None
self.joint_pos_array_seq = None
self.action_scale = None
self.stiffness_array_seq = None
self.damping_array_seq = None
joint_names = policy_config.get('joint_names')
if joint_names is None:
raise ValueError("[FSMStateWALKAMP] Missing 'joint_names' in walk_amp.yaml")
self.joint_seq = list(joint_names)
if self.action_scale_ is None:
raise ValueError("[FSMStateWALKAMP] Missing 'control.action_scale' in walk_amp.yaml")
if np.isscalar(self.action_scale_):
self.action_scale = np.full(len(self.joint_seq), float(self.action_scale_), dtype=np.float32)
else:
self.action_scale = np.array(self.action_scale_, dtype=np.float32)
if len(self.action_scale) != len(self.joint_seq):
raise ValueError(
f"[FSMStateWALKAMP] control.action_scale length {len(self.action_scale)} does not match joint count {len(self.joint_seq)}"
)
init_state_config = policy_config.get('init_state', {})
default_joint_angles = init_state_config.get('default_joint_angles')
if default_joint_angles is None:
raise ValueError("[FSMStateWALKAMP] Missing 'init_state.default_joint_angles' in walk_amp.yaml")
self.joint_pos_array_seq = np.array(default_joint_angles, dtype=np.float32)
if len(self.joint_pos_array_seq) != len(self.joint_seq):
raise ValueError(
f"[FSMStateWALKAMP] init_state.default_joint_angles length {len(self.joint_pos_array_seq)} does not match joint count {len(self.joint_seq)}"
)
gains_config = policy_config.get('gains', {})
kp_values = gains_config.get('kp')
kd_values = gains_config.get('kd')
if kp_values is None or kd_values is None:
raise ValueError("[FSMStateWALKAMP] Missing 'gains.kp' or 'gains.kd' in walk_amp.yaml")
self.stiffness_array_seq = np.array(kp_values, dtype=np.float32)
self.damping_array_seq = np.array(kd_values, dtype=np.float32)
if len(self.stiffness_array_seq) != len(self.joint_seq):
raise ValueError(
f"[FSMStateWALKAMP] gains.kp length {len(self.stiffness_array_seq)} does not match joint count {len(self.joint_seq)}"
)
if len(self.damping_array_seq) != len(self.joint_seq):
raise ValueError(
f"[FSMStateWALKAMP] gains.kd length {len(self.damping_array_seq)} does not match joint count {len(self.joint_seq)}"
)
# # 设置从序列到实验室顺序的映射
self.joint_xml = [
"hip_pitch_l_joint", "hip_roll_l_joint", "hip_yaw_l_joint",
"knee_pitch_l_joint", "ankle_pitch_l_joint", "ankle_roll_l_joint",
"hip_pitch_r_joint", "hip_roll_r_joint", "hip_yaw_r_joint",
"knee_pitch_r_joint", "ankle_pitch_r_joint", "ankle_roll_r_joint",
"waist_yaw_joint", "waist_roll_joint", "waist_pitch_joint",
"shoulder_pitch_l_joint", "shoulder_roll_l_joint", "shoulder_yaw_l_joint",
"elbow_pitch_l_joint", "elbow_yaw_l_joint", "wrist_pitch_l_joint", "wrist_roll_l_joint",
"shoulder_pitch_r_joint", "shoulder_roll_r_joint", "shoulder_yaw_r_joint",
"elbow_pitch_r_joint", "elbow_yaw_r_joint", "wrist_pitch_r_joint", "wrist_roll_r_joint",
]
# 从MjXUML顺序映射到实验室顺序
# self.mj2lab = np.array([self.joint_xml.index(joint) for joint in self.joint_seq])
self.lab2mj = []
for name in self.joint_seq:
if name not in self.joint_xml:
raise ValueError(f"[FSMStateWALKAMP] joint '{name}' from walk_amp.yaml not found in joint_xml!")
self.lab2mj.append(self.joint_xml.index(name))
self.lab2mj = np.array(self.lab2mj, dtype=int)
# 从实验室顺序映射到MjXUML顺序
# ====== 把 23 个 lab 关节 scatter 到 29 个 xml 里，多的 6 个保持默认 ======
n_mj = len(self.joint_xml)
# 29 长度，mujoco XML 顺序，先全 0 或者你想要的默认值
self.joint_pos_array = np.zeros(n_mj, dtype=np.float32)
self.stiffness_array = np.zeros(n_mj, dtype=np.float32)
self.damping_array = np.zeros(n_mj, dtype=np.float32)
# joint_pos_array_seq / stiffness_array_seq / damping_array_seq 是 23 长度，lab 顺序
for lab_idx, mj_idx in enumerate(self.lab2mj):
self.joint_pos_array[mj_idx] = self.joint_pos_array_seq[lab_idx]
self.stiffness_array[mj_idx] = self.stiffness_array_seq[lab_idx]
self.damping_array[mj_idx] = self.damping_array_seq[lab_idx]
# 设置其他参数
self.kps_lab = self.stiffness_array_seq
self.kds_lab = self.damping_array_seq
self.default_angles_lab = self.joint_pos_array_seq
self.action_scale_lab = self.action_scale
self.filtered_x_speed = 0
def _init_onnx_session(self):
"""初始化ONNX推理会话"""
try:
# 配置SessionOptions
options = ort.SessionOptions()
# 启用图优化，使用所有可用的优化（包括算子融合等）
options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
# 设置执行模式（可选，默认执行模式是顺序执行，但图优化会改变计算图）
# 设置线程数（根据CPU核心数调整）
# 建议设置为CPU物理核心数（非超线程数），因为超线程可能不会带来线性提升
options.intra_op_num_threads = 1 # 设置计算图中的运算符内部并行线程数
options.inter_op_num_threads = 1 # 设置多个运算符之间的并行线程数（如果模型有多个分支）
# 启用内存优化（避免重复分配内存）
options.enable_mem_pattern = False # 对于固定输入大小，可以设为False以避免内存规划的开销
options.enable_mem_reuse = True # 启用内存重用机制
self.ort_session_ = ort.InferenceSession(self.model_path, options, providers=['CPUExecutionProvider'])
print(f"[{self.log_prefix}-ONNX] ONNX model loaded successfully: {self.model_path}")
except Exception as e:
print(f"[{self.log_prefix}] Failed to load ONNX model: {e}")
self.ort_session_ = None
def on_enter(self):
"""进入WALKAMP状态"""
self._reset_internal_state()
print(f"[{self.log_prefix}] enter")
self.is_first_obs_ = True
self.is_first_action_ = True
self._warmup_inference_counter = 0
self.timer_gait_ = 0.0
if self.robot_data_ is not None:
try:
self._warm_start_pose = self.robot_data_.get_joint_pos().copy()
except Exception:
self._warm_start_pose.fill(0.0)
else:
self._warm_start_pose.fill(0.0)
def run(self, flag: ControlFlag):
"""运行WALKAMP状态 - 与C++版本完全一致"""
print(f"[{self.log_prefix}] run")
# Only run policy inference every decimation_ steps
gait = gait_phase(
self.timer_gait_,
self.gait_cycle,
self.left_theta_offset,
self.right_theta_offset,
self.left_phase_ratio,
self.right_phase_ratio,
).astype(np.float32)
if int(self.robot_data_.time_now_ / self.dt_) % self.decimation_ == 0:
# print(f"[FSMStateWALKAMP] Gait phase: {gait}")
self.compute_observation(flag,gait)
self.compute_actions()
# lab 顺序目标角 23 维
target_dof_pos_lab = self.actions_ * self.action_scale_lab + self.default_angles_lab
# 拿一份当前 mj 顺序的关节角（或你原来用的 default 也行）
target_dof_pos_mj = self.robot_data_.get_joint_pos().copy()
# 只更新 23 个受控 DOF
target_dof_pos_mj[self.lab2mj] = target_dof_pos_lab
commanded_pos = target_dof_pos_mj
if self._warm_start_steps > 0 and self._warmup_inference_counter < self._warm_start_steps:
self._warmup_inference_counter += 1
blend = self._warmup_inference_counter / float(self._warm_start_steps)
commanded_pos = (1.0 - blend) * self._warm_start_pose + blend * target_dof_pos_mj
base = self.robot_data_.q_d_.shape[0] - self.motor_num_
self.robot_data_.q_d_[base:base + len(self.joint_xml)] = commanded_pos
self.robot_data_.q_dot_d_[base:base + len(self.joint_xml)] = 0.0
self.robot_data_.tau_d_[base:base + len(self.joint_xml)] = 0.0
self.last_actions_[:] = self.actions_
self.timer_gait_ += self.dt_
self.robot_data_.joint_kp_p_[:len(self.joint_xml)] = self.stiffness_array
self.robot_data_.joint_kd_p_[:len(self.joint_xml)] = self.damping_array
def compute_observation(self, flag: ControlFlag, gait):
roll, pitch, yaw = (
float(self.robot_data_.imu_data_[2]),
float(self.robot_data_.imu_data_[1]),
float(self.robot_data_.imu_data_[0]),
)
quat_wxyz = self.euler_to_quaternion_scipy(roll, pitch, yaw)
q_xyzw = np.array([quat_wxyz[1], quat_wxyz[2], quat_wxyz[3], quat_wxyz[0]], dtype=np.float32)
gravity_init = self.quat_rotate_inverse_numpy(q_xyzw, np.array([0.,0.,-1.], dtype=np.float32))
x_speed_command, y_speed_command, yaw_speed_command = self.robot_data_.get_walk_cmd()
new_filtered_x_speed = 1 * x_speed_command + (1 - 1) * self.filtered_x_speed
change = new_filtered_x_speed - self.filtered_x_speed
change = np.clip(change, -0.005, 0.005)
self.filtered_x_speed = self.filtered_x_speed + change
command = np.concatenate([
np.array([
x_speed_command,
y_speed_command,
yaw_speed_command,
], dtype=np.float32),
])
print(f'Input command: {command}')
gyro = np.array([
self.robot_data_.imu_data_[3],
self.robot_data_.imu_data_[4],
self.robot_data_.imu_data_[5]
], dtype=np.float32) * self.ang_vel_scale_
q_mj = self.robot_data_.get_joint_pos()
qdot_mj = self.robot_data_.get_joint_vel()
ang_vel = self.robot_data_.get_angular_velocity()
q_mj = self.robot_data_.get_joint_pos() # mj 顺序，长度 29
dq_mj = self.robot_data_.get_joint_vel()
# 只取 23 个受控关节，变成 lab 顺序
qj = q_mj[self.lab2mj]
dqj = dq_mj[self.lab2mj]
qj = qj - self.default_angles_lab
# Observation = ang_vel(3) + gravity(3) + command(9) + q(23) + dq(23) + action(23) = 84
proprio = np.concatenate([
ang_vel , # 3 elements
gravity_init,
command,
qj,
dqj,
self.last_actions_,
gait
])
# History buffer management exactly like C++
if self.is_first_obs_:
for i in range(self.num_hist_):
start_idx = i * self.obs_size_
end_idx = start_idx + self.obs_size_
self.proprio_hist_buf_[start_idx:end_idx] = proprio
self.is_first_obs_ = False
else:
# Shift history: head((num_hist-1)*obs_size) = tail((num_hist-1)*obs_size)
shift_size = (self.num_hist_ - 1) * self.obs_size_
self.proprio_hist_buf_[:shift_size] = self.proprio_hist_buf_[self.obs_size_:]
self.proprio_hist_buf_[shift_size:] = proprio
# Clip observations exactly like C++
self.observations_ = np.clip(self.proprio_hist_buf_, -self.clip_obs_, self.clip_obs_)
@staticmethod
def euler_to_quaternion_scipy(roll, pitch, yaw, degrees=False):
r = Rotation.from_euler('xyz', [roll, pitch, yaw], degrees=degrees)
q_xyzw = r.as_quat()
return np.array([q_xyzw[3], q_xyzw[0], q_xyzw[1], q_xyzw[2]], dtype=np.float32)
@staticmethod
def quat_rotate_inverse_numpy(q_xyzw, v):
q_w = q_xyzw[3]
q_v = q_xyzw[:3]
a = v * (2.0 * q_w * q_w - 1.0)
b = np.cross(q_v, v) * (2.0 * q_w)
c = q_v * (2.0 * np.dot(q_v, v))
return a - b + c
def compute_actions(self):
if self.ort_session_ is None:
return
try:
# Prepare input tensor
input_data = self.observations_.reshape(1, -1).astype(np.float32)
# ONNX inference
input_name = self.ort_session_.get_inputs()[0].name
outputs = self.ort_session_.run(None, {input_name: input_data})
# Extract and clip actions exactly like C++
output_data = outputs[0][0]
for i in range(self.action_num_):
self.actions_[i] = np.clip(output_data[i], -self.clip_act_, self.clip_act_)
if self.is_first_action_:
print(f"[{self.log_prefix}-ONNX] First Observation:")
for i in range(self.obs_size_):
print(f"{self.observations_[i]:.6f} ", end="")
print()
self.is_first_action_ = False
except Exception as e:
print(f"[{self.log_prefix}] ONNX Runtime inference error: {e}")
def on_exit(self):
"""退出WALKAMP状态"""
print(f"[{self.log_prefix}] exit")
# 关掉 obs 日志文件
if getattr(self, "obs_log_file", None) is not None:
try:
self.obs_log_file.flush()
self.obs_log_file.close()
print(f"[{self.log_prefix}] obs log saved to {self.obs_log_path}")
except Exception as e:
print(f"[{self.log_prefix}] failed to close obs log: {e}")
self.obs_log_file = None
def check_transition(self, flag: ControlFlag) -> FSMStateName:
"""检查状态转换"""
if flag.fsm_state_command == "gotoSTOP":
return FSMStateName.STOP
elif flag.fsm_state_command == "gotoWALKAMP":
return FSMStateName.WALKAMP
elif flag.fsm_state_command == "gotoXSIMRUN":
return FSMStateName.XSIMRUN
elif flag.fsm_state_command == "gotoZERO":
return FSMStateName.ZERO
else:
return None # 无状态转换

BIN
Deploy_Tienkung/policy/walk_amp/model/policy.onnx Normal file
View File

Binary file not shown.