362 lines
15 KiB
Python
362 lines
15 KiB
Python
#!/usr/bin/env python3
|
|
"""ROS 2 MuJoCo bridge for the original M20 and the 8-DoF FireAnt model."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import mujoco
|
|
import mujoco.viewer
|
|
import numpy as np
|
|
import rclpy
|
|
from builtin_interfaces.msg import Time
|
|
from drdds.msg import (
|
|
ImuData,
|
|
ImuDataValue,
|
|
JointData,
|
|
JointsData,
|
|
JointsDataCmd,
|
|
JointsDataValue,
|
|
MetaType,
|
|
)
|
|
from rclpy.node import Node
|
|
|
|
|
|
CURRENT_DIR = Path(__file__).resolve().parent
|
|
PROJECT_ROOT = CURRENT_DIR.parents[2]
|
|
MESSAGE_JOINT_COUNT = 16
|
|
DT = 0.001
|
|
STATE_PUBLISH_DECIMATION = 5 # 200 Hz, matching the C++ controller loop.
|
|
|
|
FIRE_ANT_JOINTS = (
|
|
"fl_knee_joint", "fl_wheel_joint",
|
|
"fr_knee_joint", "fr_wheel_joint",
|
|
"hl_knee_joint", "hl_wheel_joint",
|
|
"hr_knee_joint", "hr_wheel_joint",
|
|
)
|
|
FIRE_ANT_HOME = np.array([0.5, 0.0, 0.5, 0.0, -0.5, 0.0, -0.5, 0.0], dtype=np.float64)
|
|
|
|
M20_DIR = np.array(
|
|
[1, 1, -1, 1, 1, -1, 1, -1, -1, 1, -1, 1, -1, -1, 1, -1],
|
|
dtype=np.float64,
|
|
)
|
|
M20_OFFSET = np.deg2rad(
|
|
np.array([-25, 229, 160, 0, 25, -131, -200, 0, -25, -229, -160, 0, 25, 131, 200, 0])
|
|
)
|
|
M20_HOME = np.array(
|
|
[-0.438, -1.16, 2.76, 0, 0.438, -1.16, 2.76, 0,
|
|
-0.438, 1.16, -2.76, 0, 0.438, 1.16, -2.76, 0],
|
|
dtype=np.float64,
|
|
)
|
|
|
|
|
|
def model_path(robot: str, terrain: str) -> Path:
|
|
if robot == "fire_ant":
|
|
name = "fire_ant_rough.xml" if terrain == "rough" else "fire_ant.xml"
|
|
return PROJECT_ROOT / "M20_description" / "fire_ant_mjcf" / "mjcf" / name
|
|
name = "M20_stair.xml" if terrain == "rough" else "M20.xml"
|
|
return PROJECT_ROOT / "M20_description" / "m20_mjcf" / "mjcf" / name
|
|
|
|
|
|
def quaternion_to_euler(q: np.ndarray) -> np.ndarray:
|
|
"""Convert a MuJoCo wxyz quaternion to roll, pitch and yaw."""
|
|
w, x, y, z = q
|
|
roll = np.arctan2(2.0 * (w * x + y * z), 1.0 - 2.0 * (x * x + y * y))
|
|
pitch = np.arcsin(np.clip(2.0 * (w * y - z * x), -1.0, 1.0))
|
|
yaw = np.arctan2(2.0 * (w * z + x * y), 1.0 - 2.0 * (y * y + z * z))
|
|
return np.array([roll, pitch, yaw], dtype=np.float64)
|
|
|
|
|
|
class MuJoCoSimulationNode(Node):
|
|
def __init__(
|
|
self,
|
|
robot: str,
|
|
xml_path: Path,
|
|
use_viewer: bool,
|
|
randomize: bool,
|
|
seed: int,
|
|
realtime_factor: float,
|
|
) -> None:
|
|
super().__init__("mujoco_simulation")
|
|
if not xml_path.is_file():
|
|
raise FileNotFoundError(f"Cannot find MJCF: {xml_path}")
|
|
|
|
self.robot = robot
|
|
self.model = mujoco.MjModel.from_xml_path(str(xml_path))
|
|
self.model.opt.timestep = DT
|
|
self.data = mujoco.MjData(self.model)
|
|
self.realtime_factor = realtime_factor
|
|
|
|
self.dof_num = self.model.nu
|
|
expected_dof = 8 if robot == "fire_ant" else 16
|
|
if self.dof_num != expected_dof:
|
|
raise RuntimeError(f"{robot} expected {expected_dof} actuators, model has {self.dof_num}")
|
|
|
|
self.joint_names = list(FIRE_ANT_JOINTS) if robot == "fire_ant" else [
|
|
mujoco.mj_id2name(
|
|
self.model,
|
|
mujoco.mjtObj.mjOBJ_JOINT,
|
|
int(self.model.actuator_trnid[i, 0]),
|
|
)
|
|
for i in range(self.dof_num)
|
|
]
|
|
self.joint_ids = np.array(
|
|
[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, name) for name in self.joint_names],
|
|
dtype=np.int32,
|
|
)
|
|
if np.any(self.joint_ids < 0):
|
|
raise RuntimeError(f"Missing joints in MJCF: {self.joint_names}")
|
|
self.qpos_adr = self.model.jnt_qposadr[self.joint_ids]
|
|
self.dof_adr = self.model.jnt_dofadr[self.joint_ids]
|
|
self.actuator_ids = np.array(
|
|
[
|
|
next(i for i in range(self.model.nu) if int(self.model.actuator_trnid[i, 0]) == joint_id)
|
|
for joint_id in self.joint_ids
|
|
],
|
|
dtype=np.int32,
|
|
)
|
|
|
|
self.joint_dir = np.ones(self.dof_num) if robot == "fire_ant" else M20_DIR.copy()
|
|
self.pos_offset = np.zeros(self.dof_num) if robot == "fire_ant" else M20_OFFSET.copy()
|
|
self.kp_cmd = np.zeros(self.dof_num)
|
|
self.kd_cmd = np.zeros(self.dof_num)
|
|
self.pos_cmd = np.zeros(self.dof_num)
|
|
self.vel_cmd = np.zeros(self.dof_num)
|
|
self.tau_ff = np.zeros(self.dof_num)
|
|
self.applied_tau = np.zeros(self.dof_num)
|
|
self.velocity_limit = (
|
|
np.array([17.8024, 109.956] * 4)
|
|
if robot == "fire_ant"
|
|
else np.array([45.0, 22.4, 22.4, 100.0] * 4)
|
|
)
|
|
self.kp_scale = np.ones(self.dof_num)
|
|
self.kd_scale = np.ones(self.dof_num)
|
|
self.timestamp = 0.0
|
|
|
|
self.base_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, "base_link")
|
|
self.wheel_body_ids: list[int] = []
|
|
self.knee_body_ids: list[int] = []
|
|
if robot == "fire_ant":
|
|
self.wheel_body_ids = [self._body_id(f"{leg}_wheel") for leg in ("fl", "fr", "hl", "hr")]
|
|
self.knee_body_ids = [self._body_id(f"{leg}_knee") for leg in ("fl", "fr", "hl", "hr")]
|
|
if randomize:
|
|
self._apply_training_randomization(seed)
|
|
|
|
self._set_initial_pose()
|
|
self.imu_pub = self.create_publisher(ImuData, "/IMU_DATA", 200)
|
|
self.joints_pub = self.create_publisher(JointsData, "/JOINTS_DATA", 200)
|
|
self.cmd_sub = self.create_subscription(JointsDataCmd, "/JOINTS_CMD", self._cmd_callback, 50)
|
|
self.viewer = mujoco.viewer.launch_passive(self.model, self.data) if use_viewer else None
|
|
|
|
mode = "randomized" if randomize else "nominal"
|
|
self.get_logger().info(
|
|
f"Loaded {robot} ({mode}) from {xml_path}; dof={self.dof_num}, dt={DT:.4f}s"
|
|
)
|
|
|
|
def _body_id(self, name: str) -> int:
|
|
body_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, name)
|
|
if body_id < 0:
|
|
raise RuntimeError(f"Missing body in MJCF: {name}")
|
|
return body_id
|
|
|
|
def _apply_training_randomization(self, seed: int) -> None:
|
|
"""Apply one reproducible sample from the training-time domain ranges."""
|
|
rng = np.random.default_rng(seed)
|
|
original_mass = self.model.body_mass.copy()
|
|
for body_id in [*self.knee_body_ids, *self.wheel_body_ids]:
|
|
mass_scale = rng.uniform(0.85, 1.15)
|
|
inertia_scale = mass_scale * rng.uniform(0.85, 1.15)
|
|
self.model.body_mass[body_id] *= mass_scale
|
|
self.model.body_inertia[body_id] *= inertia_scale
|
|
|
|
base_mass = max(0.1, original_mass[self.base_id] + rng.uniform(-1.0, 3.0))
|
|
base_mass_scale = base_mass / original_mass[self.base_id]
|
|
self.model.body_mass[self.base_id] = base_mass
|
|
self.model.body_inertia[self.base_id] *= base_mass_scale * rng.uniform(0.85, 1.15)
|
|
self.model.body_ipos[self.base_id] += rng.uniform((-0.03, -0.03, -0.02), (0.03, 0.03, 0.02))
|
|
|
|
for body_id in self.wheel_body_ids:
|
|
geom_start = self.model.body_geomadr[body_id]
|
|
geom_count = self.model.body_geomnum[body_id]
|
|
self.model.geom_friction[geom_start:geom_start + geom_count, 0] = rng.uniform(0.3, 1.0)
|
|
self.kp_scale = rng.uniform(0.85, 1.15, self.dof_num)
|
|
self.kd_scale = rng.uniform(0.85, 1.15, self.dof_num)
|
|
|
|
def _set_initial_pose(self) -> None:
|
|
self.data.qpos[:] = 0.0
|
|
self.data.qvel[:] = 0.0
|
|
self.data.qpos[0:3] = (0.0, 0.0, 0.306 if self.robot == "fire_ant" else 0.2)
|
|
self.data.qpos[3:7] = (1.0, 0.0, 0.0, 0.0)
|
|
self.data.qpos[self.qpos_adr] = FIRE_ANT_HOME if self.robot == "fire_ant" else M20_HOME
|
|
mujoco.mj_forward(self.model, self.data)
|
|
|
|
def _cmd_callback(self, msg: JointsDataCmd) -> None:
|
|
if len(msg.data.joints_data) < self.dof_num:
|
|
self.get_logger().warn(
|
|
f"Ignoring short JointsDataCmd ({len(msg.data.joints_data)} < {self.dof_num})"
|
|
)
|
|
return
|
|
pub_pos = np.empty(self.dof_num)
|
|
pub_vel = np.empty(self.dof_num)
|
|
for i in range(self.dof_num):
|
|
command = msg.data.joints_data[i]
|
|
self.kp_cmd[i] = command.kp
|
|
self.kd_cmd[i] = command.kd
|
|
pub_pos[i] = command.position
|
|
pub_vel[i] = command.velocity
|
|
self.tau_ff[i] = command.torque
|
|
self.pos_cmd = pub_pos * self.joint_dir + self.pos_offset
|
|
self.vel_cmd = pub_vel * self.joint_dir
|
|
|
|
def _joint_state(self) -> tuple[np.ndarray, np.ndarray]:
|
|
return self.data.qpos[self.qpos_adr].copy(), self.data.qvel[self.dof_adr].copy()
|
|
|
|
def _apply_joint_torque(self) -> None:
|
|
q, dq = self._joint_state()
|
|
torque = (
|
|
self.kp_cmd * self.kp_scale * (self.pos_cmd - q)
|
|
+ self.kd_cmd * self.kd_scale * (self.vel_cmd - dq)
|
|
+ self.tau_ff
|
|
)
|
|
ctrl_range = self.model.actuator_ctrlrange[self.actuator_ids]
|
|
torque = np.clip(torque, ctrl_range[:, 0], ctrl_range[:, 1])
|
|
# The MuJoCo motor has a torque limit but no built-in speed limit. Prevent
|
|
# positive work beyond the datasheet/training actuator velocity limit.
|
|
overspeed_accelerating = (np.abs(dq) >= self.velocity_limit) & (torque * dq > 0.0)
|
|
torque[overspeed_accelerating] = 0.0
|
|
self.data.ctrl[self.actuator_ids] = torque
|
|
|
|
def _body_linear_velocity(self) -> np.ndarray:
|
|
# Free-joint translational qvel is in world axes; the policy uses body axes.
|
|
rotation_world_from_body = self.data.xmat[self.base_id].reshape(3, 3)
|
|
return rotation_world_from_body.T @ self.data.qvel[0:3]
|
|
|
|
def _contact_force_norms(self, body_ids: list[int]) -> np.ndarray:
|
|
# cfrc_ext stores [torque, force]; norm is frame invariant.
|
|
return np.linalg.norm(self.data.cfrc_ext[body_ids, 3:6], axis=1)
|
|
|
|
def _make_stamp(self) -> Time:
|
|
stamp = Time()
|
|
stamp.sec = int(self.timestamp)
|
|
stamp.nanosec = int((self.timestamp - stamp.sec) * 1e9)
|
|
return stamp
|
|
|
|
def _publish_robot_state(self) -> None:
|
|
q_world = self.data.sensordata[0:4]
|
|
rpy_deg = np.rad2deg(quaternion_to_euler(q_world))
|
|
body_acc = self.data.sensordata[4:7]
|
|
angular_velocity_body = self.data.sensordata[7:10]
|
|
|
|
imu_msg = ImuData()
|
|
imu_msg.header = MetaType()
|
|
imu_msg.header.frame_id = 0
|
|
imu_msg.header.stamp = self._make_stamp()
|
|
imu_msg.data = ImuDataValue()
|
|
imu_msg.data.roll, imu_msg.data.pitch, imu_msg.data.yaw = map(float, rpy_deg)
|
|
imu_msg.data.omega_x, imu_msg.data.omega_y, imu_msg.data.omega_z = map(
|
|
float, angular_velocity_body
|
|
)
|
|
imu_msg.data.acc_x, imu_msg.data.acc_y, imu_msg.data.acc_z = map(float, body_acc)
|
|
self.imu_pub.publish(imu_msg)
|
|
|
|
q, dq = self._joint_state()
|
|
pub_pos = (q - self.pos_offset) * self.joint_dir
|
|
pub_vel = dq * self.joint_dir
|
|
self.applied_tau = self.data.actuator_force[self.actuator_ids] * self.joint_dir
|
|
|
|
joints_msg = JointsData()
|
|
joints_msg.header = MetaType()
|
|
joints_msg.header.frame_id = 0
|
|
joints_msg.header.stamp = self._make_stamp()
|
|
joints_msg.data = JointsDataValue()
|
|
joints_msg.data.joints_data = [JointData() for _ in range(MESSAGE_JOINT_COUNT)]
|
|
for i in range(self.dof_num):
|
|
joint = joints_msg.data.joints_data[i]
|
|
joint.name = [32, 32, 32, 32]
|
|
joint.data_id = 0
|
|
joint.status_word = 1
|
|
joint.position = float(pub_pos[i])
|
|
joint.velocity = float(pub_vel[i])
|
|
joint.torque = float(self.applied_tau[i])
|
|
joint.motion_temp = 40.0
|
|
joint.driver_temp = 45.0
|
|
|
|
if self.robot == "fire_ant":
|
|
# Sim-only extension inside the eight unused fields of the fixed 16-joint DDS message.
|
|
wheel_force = self._contact_force_norms(self.wheel_body_ids)
|
|
knee_force = self._contact_force_norms(self.knee_body_ids)
|
|
base_force = self._contact_force_norms([self.base_id])[0]
|
|
base_velocity = self._body_linear_velocity()
|
|
for i in range(4):
|
|
joints_msg.data.joints_data[8 + i].position = float(wheel_force[i])
|
|
joints_msg.data.joints_data[12 + i].position = float(knee_force[i])
|
|
for i in range(3):
|
|
joints_msg.data.joints_data[8 + i].torque = float(base_velocity[i])
|
|
joints_msg.data.joints_data[12].velocity = float(base_force)
|
|
|
|
self.joints_pub.publish(joints_msg)
|
|
|
|
def start(self) -> None:
|
|
step = 0
|
|
next_step_time = time.perf_counter()
|
|
while rclpy.ok() and (self.viewer is None or self.viewer.is_running()):
|
|
rclpy.spin_once(self, timeout_sec=0.0)
|
|
self._apply_joint_torque()
|
|
mujoco.mj_step(self.model, self.data)
|
|
mujoco.mj_rnePostConstraint(self.model, self.data)
|
|
step += 1
|
|
self.timestamp = step * DT
|
|
if step % STATE_PUBLISH_DECIMATION == 0:
|
|
self._publish_robot_state()
|
|
if self.viewer is not None and step % 20 == 0:
|
|
self.viewer.sync()
|
|
|
|
if self.realtime_factor > 0.0:
|
|
next_step_time += DT / self.realtime_factor
|
|
remaining = next_step_time - time.perf_counter()
|
|
if remaining > 0:
|
|
time.sleep(remaining)
|
|
elif remaining < -0.1:
|
|
next_step_time = time.perf_counter()
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--robot", choices=("fire_ant", "m20"), default="fire_ant")
|
|
parser.add_argument("--terrain", choices=("flat", "rough"), default="rough")
|
|
parser.add_argument("--xml", type=Path, help="Override the selected MJCF path")
|
|
parser.add_argument("--headless", action="store_true", help="Run without the MuJoCo viewer")
|
|
parser.add_argument("--randomize", action="store_true", help="Sample training mass/friction/gain ranges once")
|
|
parser.add_argument("--seed", type=int, default=0)
|
|
parser.add_argument("--realtime-factor", type=float, default=1.0, help="0 runs as fast as possible")
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
xml_path = args.xml.resolve() if args.xml else model_path(args.robot, args.terrain).resolve()
|
|
rclpy.init()
|
|
node = MuJoCoSimulationNode(
|
|
robot=args.robot,
|
|
xml_path=xml_path,
|
|
use_viewer=not args.headless,
|
|
randomize=args.randomize,
|
|
seed=args.seed,
|
|
realtime_factor=args.realtime_factor,
|
|
)
|
|
try:
|
|
node.start()
|
|
except KeyboardInterrupt:
|
|
pass
|
|
finally:
|
|
node.destroy_node()
|
|
if rclpy.ok():
|
|
rclpy.shutdown()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
np.set_printoptions(precision=4, suppress=True)
|
|
main()
|