[software] 添加16DOF早期训练仿真与Sim2Real闭环
This commit is contained in:
@@ -0,0 +1,121 @@
|
||||
import numpy as np
|
||||
import mujoco
|
||||
import matplotlib.pyplot as plt
|
||||
from .math_utils import get_body_velocity
|
||||
|
||||
class DisturbanceTester:
|
||||
def __init__(self, interval=3.0, duration=0.02, force_mag=2000.0):
|
||||
self.interval = interval
|
||||
self.duration = duration
|
||||
self.force_mag = force_mag
|
||||
self.time_log = []
|
||||
self.base_vel_log = []
|
||||
self.vel_log = []
|
||||
self.force_input_log = []
|
||||
|
||||
# Define sensors to plot
|
||||
self.SENSOR_NAMES_TO_PLOT = {
|
||||
"FR": ["FR_hip_torque", "FR_thigh_torque", "FR_calf_torque"],
|
||||
"FL": ["FL_hip_torque", "FL_thigh_torque", "FL_calf_torque"],
|
||||
"RR": ["RR_hip_torque", "RR_thigh_torque", "RR_calf_torque"],
|
||||
"RL": ["RL_hip_torque", "RL_thigh_torque", "RL_calf_torque"],
|
||||
"Wheels": ["FR_wheel_torque", "FL_wheel_torque", "RR_wheel_torque", "RL_wheel_torque"]
|
||||
}
|
||||
self.sensor_logs = {name: [] for group in self.SENSOR_NAMES_TO_PLOT.values() for name in group}
|
||||
|
||||
def update(self, current_time, m, d, estimated_vel):
|
||||
cycle_time = current_time % self.interval
|
||||
is_pushing = cycle_time < self.duration
|
||||
applied_force = np.zeros(6)
|
||||
if is_pushing:
|
||||
applied_force[0] = -self.force_mag
|
||||
|
||||
# Apply external force
|
||||
base_body_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_BODY, "base_link")
|
||||
if base_body_id == -1:
|
||||
base_body_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_BODY, "trunk")
|
||||
if base_body_id != -1:
|
||||
d.xfrc_applied[base_body_id] = applied_force
|
||||
|
||||
# Record data
|
||||
self.time_log.append(current_time)
|
||||
true_vel_body = get_body_velocity(m, d)
|
||||
self.base_vel_log.append(true_vel_body)
|
||||
self.force_input_log.append(applied_force[0])
|
||||
self.vel_log.append(estimated_vel)
|
||||
|
||||
# Record sensor data
|
||||
for name in self.sensor_logs.keys():
|
||||
sid = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_SENSOR, name)
|
||||
if sid != -1:
|
||||
adr = m.sensor_adr[sid]
|
||||
val = d.sensordata[adr]
|
||||
self.sensor_logs[name].append(val)
|
||||
else:
|
||||
self.sensor_logs[name].append(0.0)
|
||||
|
||||
def plot_results(self):
|
||||
print("Generating diagnostic plots...")
|
||||
if len(self.time_log) == 0:
|
||||
print("No data recorded. Skipping plot.")
|
||||
return
|
||||
|
||||
time_arr = np.array(self.time_log)
|
||||
vel_arr = np.array(self.base_vel_log)
|
||||
force_arr = np.array(self.force_input_log)
|
||||
|
||||
fig, axes = plt.subplots(4, 2, figsize=(16, 18), sharex=True)
|
||||
|
||||
# 1. External Force
|
||||
axes[0, 0].plot(time_arr, force_arr, 'r-', linewidth=1.5)
|
||||
axes[0, 0].set_title("External Push Force (N)")
|
||||
axes[0, 0].set_ylabel("Force")
|
||||
axes[0, 0].grid(True)
|
||||
|
||||
# 2. Velocity
|
||||
axes[0, 1].plot(time_arr, vel_arr[:, 0], label='True Vx')
|
||||
axes[0, 1].plot(time_arr, vel_arr[:, 1], label='True Vy')
|
||||
est_vel_arr = np.array(self.vel_log)
|
||||
if est_vel_arr.shape[1] >= 2:
|
||||
axes[0, 1].plot(time_arr, est_vel_arr[:, 0], '--', label='Est Vx')
|
||||
axes[0, 1].plot(time_arr, est_vel_arr[:, 1], '--', label='Est Vy')
|
||||
axes[0, 1].set_title("Base Velocity (m/s)")
|
||||
axes[0, 1].legend()
|
||||
axes[0, 1].grid(True)
|
||||
|
||||
# 3. Legs
|
||||
plot_config = [
|
||||
("FR", axes[1, 1]),
|
||||
("FL", axes[1, 0]),
|
||||
("RR", axes[2, 1]),
|
||||
("RL", axes[2, 0])
|
||||
]
|
||||
for group_name, ax in plot_config:
|
||||
sensor_names = self.SENSOR_NAMES_TO_PLOT[group_name]
|
||||
labels = ["Hip", "Thigh", "Calf"]
|
||||
for i, s_name in enumerate(sensor_names):
|
||||
if s_name in self.sensor_logs:
|
||||
data = self.sensor_logs[s_name]
|
||||
ax.plot(time_arr, data, label=labels[i], linewidth=1)
|
||||
ax.set_title(f"{group_name} Leg Torques")
|
||||
ax.set_ylabel("Torque (Nm)")
|
||||
ax.legend(loc='upper right')
|
||||
ax.grid(True, alpha=0.3)
|
||||
|
||||
# 4. Wheels
|
||||
ax_wheel = axes[3, 0]
|
||||
wheel_sensors = self.SENSOR_NAMES_TO_PLOT["Wheels"]
|
||||
for w_name in wheel_sensors:
|
||||
if w_name in self.sensor_logs:
|
||||
data = self.sensor_logs[w_name]
|
||||
short_label = w_name.replace("_wheel_torque", "")
|
||||
ax_wheel.plot(time_arr, data, label=short_label, linewidth=1)
|
||||
ax_wheel.set_title("Wheel Torques")
|
||||
ax_wheel.set_ylabel("Torque (Nm)")
|
||||
ax_wheel.set_xlabel("Time (s)")
|
||||
ax_wheel.legend()
|
||||
ax_wheel.grid(True, alpha=0.3)
|
||||
|
||||
axes[3, 1].axis('off')
|
||||
plt.tight_layout()
|
||||
plt.show()
|
||||
@@ -0,0 +1,26 @@
|
||||
import time
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
from tools.math_utils import get_body_velocity
|
||||
|
||||
class SimpleLogger:
|
||||
def __init__(self, log_dir="."):
|
||||
self.log_file = Path(log_dir) / f"sim2sim_log_{int(time.time())}.txt"
|
||||
self.data = []
|
||||
print(f"[Logger] Will write diagnostic data to: {self.log_file}")
|
||||
|
||||
def update(self, sim_time, m, d, command):
|
||||
true_vel_body = get_body_velocity(m, d)
|
||||
# 记录 仿真时间、控制指令、真实的机体线速度
|
||||
self.data.append(
|
||||
f"{sim_time:.4f}, {command[0]:.4f}, {command[1]:.4f}, {command[2]:.4f}, "
|
||||
f"{true_vel_body[0]:.4f}, {true_vel_body[1]:.4f}, {true_vel_body[2]:.4f}"
|
||||
)
|
||||
|
||||
def save(self):
|
||||
print(f"\n[Logger] Saving {len(self.data)} records to {self.log_file}...")
|
||||
with open(self.log_file, "w") as f:
|
||||
f.write("time, cmd_vx, cmd_vy, cmd_yaw, true_vx, true_vy, true_vz\n")
|
||||
for row in self.data:
|
||||
f.write(row + "\n")
|
||||
print("[Logger] Save complete.")
|
||||
@@ -0,0 +1,35 @@
|
||||
import numpy as np
|
||||
|
||||
def get_gravity_orientation(quat_wxyz):
|
||||
"""Compute projected gravity in body frame from quaternion [w,x,y,z]."""
|
||||
qw, qx, qy, qz = quat_wxyz
|
||||
gx = 2.0 * (-qz * qx + qw * qy)
|
||||
gy = -2.0 * (qz * qy + qw * qx)
|
||||
gz = 1.0 - 2.0 * (qw * qw + qz * qz)
|
||||
return np.array([gx, gy, gz], dtype=np.float32)
|
||||
|
||||
def quat_rotate_inverse(quat_wxyz, v):
|
||||
"""Rotate vector v from world frame to body frame."""
|
||||
q_w = quat_wxyz[0]
|
||||
q_vec = quat_wxyz[1:]
|
||||
a = v * (2.0 * q_w * q_w - 1.0)
|
||||
b = np.cross(q_vec, v) * q_w * 2.0
|
||||
c = q_vec * np.dot(q_vec, v) * 2.0
|
||||
return a - b + c
|
||||
|
||||
def get_body_velocity(m, d):
|
||||
"""Calculate body frame linear velocity (vx, vy, vz) from world velocity."""
|
||||
v_world = d.qvel[0:3]
|
||||
q = d.qpos[3:7] # MuJoCo freejoint quaternion is at indices 3:7 (x,y,z, w,x,y,z)
|
||||
w, x, y, z = q
|
||||
norm = np.sqrt(w*w + x*x + y*y + z*z)
|
||||
if norm < 1e-6:
|
||||
return np.array([0.0, 0.0, 0.0], dtype=np.float32)
|
||||
w, x, y, z = w/norm, x/norm, y/norm, z/norm
|
||||
R = np.array([
|
||||
[1 - 2*y*y - 2*z*z, 2*x*y - 2*z*w, 2*x*z + 2*y*w],
|
||||
[2*x*y + 2*z*w, 1 - 2*x*x - 2*z*z, 2*y*z - 2*x*w],
|
||||
[2*x*z - 2*y*w, 2*y*z + 2*x*w, 1 - 2*x*x - 2*y*y]
|
||||
])
|
||||
v_body = R.T @ v_world
|
||||
return v_body.astype(np.float32)
|
||||
Reference in New Issue
Block a user