diff --git a/01_doc/version_history.md b/01_doc/version_history.md index a45b9ac..e5a5a06 100644 --- a/01_doc/version_history.md +++ b/01_doc/version_history.md @@ -11,6 +11,7 @@ | `v0.4.0` | 新训练基线 | 第一份完整的新 MJCF、新 mjlab 框架和 Rough 策略工程 | | `v0.5.0` | 随机化增强 | 扩大观测、延迟和动力学随机化，加入持续外力扰动 | | `v0.6.0` | 比赛训练架构 | 分轴奖励、自适应指令课程、障碍释放课程和比赛站姿 | +| `v0.7.0` | MuJoCo 工具 | 姿态优化、IK 扫描、动力学、MPC 和 GUI 调试工具 | ## `v0.4.0` 的模型变化 @@ -42,3 +43,11 @@ - 该 Tag 保存比赛训练代码架构，不把每次继续训练产生的 checkpoint 误记为新的软件版本。 训练阶段的详细关系见 [`training_evolution.md`](training_evolution.md)。 + +## `v0.7.0` 的独立 MuJoCo 工具 + +- 比赛训练架构、MJCF、模型和依赖锁文件保持 `v0.6.0` 状态不变。 +- 增加解析姿态表、RL 友好姿态筛选和 MuJoCo 静态姿态优化。 +- 增加 IK/差速轮参数扫描，可导出 JSON 结果。 +- 增加 Robot、Controller、Dynamics、MPCController 和 GUI 调试链路。 +- 记录历史工具常量与新版 MJCF 质量、比赛默认站姿之间的参数边界，避免将分析结果直接当作已校准真机参数。 diff --git a/05_software/README.md b/05_software/README.md index 9c0299a..3ed6053 100644 --- a/05_software/README.md +++ b/05_software/README.md @@ -5,7 +5,7 @@ ```text 05_software/ ├─ train/ -│ └─ rc_mjlab/ # 训练、MJCF、Sim2Sim 和本地 mjlab 依赖 +│ └─ rc_mjlab/ # 训练、MJCF、MuJoCo、Sim2Sim 和本地 mjlab 依赖 └─ real/ ├─ ik_real/ # IK 轨迹与早期真机控制 └─ sim2real/ # 第一代 Python 策略真机部署 @@ -18,6 +18,8 @@ MJCF + mjlab task | v PPO 训练策略 + | + +----> MuJoCo 姿态 / IK / MPC 调试 | +----> Sim2Sim 策略验证 | @@ -26,7 +28,7 @@ MJCF + mjlab task IK real --------------------------------> 电机 ``` -`rc_mjlab` 是自包含工程。训练、MJCF、Sim2Sim 和策略权重通过相对路径绑定，因此保留其内部布局，没有为了目录外观拆散。第一代完整闭环见 `v0.3.0`，第一份新版 MJCF 与训练框架见 `v0.4.0`，随机化增强版见 `v0.5.0`，比赛最终训练架构见 `v0.6.0`。 +`rc_mjlab` 是自包含工程。训练、MJCF、MuJoCo、Sim2Sim 和策略权重通过相对路径绑定，因此保留其内部布局，没有为了目录外观拆散。第一代完整闭环见 `v0.3.0`，第一份新版 MJCF 与训练框架见 `v0.4.0`，随机化增强版见 `v0.5.0`，比赛最终训练架构见 `v0.6.0`，后期 MuJoCo 工具集见 `v0.7.0`。 详细说明见： diff --git a/05_software/train/README.md b/05_software/train/README.md index f703ece..fbb3db4 100644 --- a/05_software/train/README.md +++ b/05_software/train/README.md @@ -2,13 +2,14 @@ `rc_mjlab/` 保存 16DOF 轮足机器人的当前训练与 Sim2Sim 工程。历史快照由 Git Tag 保留，不在目录中复制 `old`、`new` 或 `final` 版本。 -当前内容对应 `v0.6.0`，是比赛使用的最终训练代码架构。训练过程可能先获得基模，再调整奖励、课程和环境参数继续训练；模型 checkpoint 的变化不等同于软件架构变化。 +当前内容对应 `v0.7.0`：训练代码保持 `v0.6.0` 的比赛架构，新增后期 MuJoCo 姿态、IK、动力学和 MPC 工具。训练过程可能先获得基模，再调整奖励、课程和环境参数继续训练；模型 checkpoint 的变化不等同于软件架构变化。 ## 内容 - `src/robot`：Flat、Rough、Crawl 训练任务和自定义 MDP - `mjcf`：轮足机器人 MuJoCo 模型和网格 - `sim2sim`：策略加载、交互控制和比赛地形验证 +- `mujoco_sim`：不依赖训练循环的姿态、IK、动力学和 MPC 分析 - `mjlab`：固定版本的本地训练框架依赖 - `model_rough.pt`：本阶段 Rough 策略权重 - `pyproject.toml`、`uv.lock`：Python 环境与依赖锁定 @@ -19,4 +20,6 @@ `v0.6.0` 在 `v0.5.0` 之后转向比赛任务优化：降低部分过强随机化，加入分轴速度跟踪奖励、自适应指令课程、障碍地形释放课程、楼梯横向/偏航约束以及更完整的训练诊断。详细对比见 [`../../01_doc/training_evolution.md`](../../01_doc/training_evolution.md)。 +`v0.7.0` 不修改比赛训练架构，增加独立 MuJoCo 工具；入口和参数边界见 [`rc_mjlab/mujoco_sim/README.md`](rc_mjlab/mujoco_sim/README.md)。 + 工程命令和任务说明见 [`rc_mjlab/README.md`](rc_mjlab/README.md)，本地依赖来源见 [`rc_mjlab/DEPENDENCIES.md`](rc_mjlab/DEPENDENCIES.md)。 diff --git a/05_software/train/rc_mjlab/README.md b/05_software/train/rc_mjlab/README.md index 55bcc4c..283d13f 100644 --- a/05_software/train/rc_mjlab/README.md +++ b/05_software/train/rc_mjlab/README.md @@ -2,7 +2,7 @@ 基于 [mjlab](https://github.com/google-deepmind/mjlab) 框架的四轮腿混合机器人强化学习训练与部署部署项目，面向机器人竞赛场景（如越障、匍匐、斜坡、台阶等复合任务）。 -> 本目录对应 `v0.6.0`：比赛使用的最终训练架构。它在前两版训练代码上重新平衡随机化强度，引入分轴速度奖励、自适应命令课程、障碍逐步释放、楼梯稳定约束和训练诊断指标。当前目录中的 `model_rough.pt` 是早期参考权重；比赛最终使用的 `model_6800.onnx` 将随最终部署版本归档。 +> 当前目录对应 `v0.7.0`：保留 `v0.6.0` 的比赛训练架构，并加入后期 MuJoCo 独立工具集。当前目录中的 `model_rough.pt` 是早期参考权重；比赛最终使用的 `model_6800.onnx` 将随最终部署版本归档。 --- @@ -53,6 +53,7 @@ rc_mjlab/ │ ├── wheelleg.xml # 机器人 MuJoCo 模型（含网格引用） │ ├── scene.xml # mjlab 场景入口文件 │ └── meshes/ # STL/OBJ 碰撞与外观网格 +├── mujoco_sim/ # 姿态、IK、动力学和 MPC 独立工具 ├── model_rough.pt # 本阶段用于回放和 Sim2Sim 的 Rough 策略 ├── pyproject.toml # 项目依赖（uv 管理，含清华镜像源加速） └── uv.lock # 精确依赖锁定文件 diff --git a/05_software/train/rc_mjlab/mujoco_sim/README.md b/05_software/train/rc_mjlab/mujoco_sim/README.md new file mode 100644 index 0000000..82172b3 --- /dev/null +++ b/05_software/train/rc_mjlab/mujoco_sim/README.md @@ -0,0 +1,57 @@ +# MuJoCo 独立工具集 + +本目录保存不依赖训练循环的 MuJoCo、姿态、IK、动力学和 MPC 分析工具。脚本通过父目录读取 `../mjcf/scene.xml` 与 `../mjcf/wheelleg.xml`，因此应从 `rc_mjlab` 工程根目录运行。 + +## 工具分类 + +| 入口 | 用途 | +| --- | --- | +| `posture_tool.py` | 基于解析运动学快速生成站立和低姿态参数表 | +| `rl_friendly_opt.py` | 按轮心位置、关节力矩和雅可比条件数筛选适合 RL 的姿态 | +| `posture_optimizer.py` | 解析计算与 MuJoCo 扫描结合的姿态优化 | +| `static_posture_optimizer.py` | 在重力和地面接触下评估静态站姿、支撑域和离地间隙 | +| `ik_diff_sweep.py` | 扫描 IK 姿态和差速轮跟踪参数，可导出 JSON | +| `run.py` | 启动完整 MuJoCo 控制、GUI 和 MPC 调试链路 | +| `robot.py`、`controller.py` | 仿真机器人接口和控制器 | +| `dynamics.py`、`mpc.py`、`mpc_controller.py` | Pinocchio 动力学与 OSQP MPC | + +## 依赖 + +执行工程根目录的 `uv sync` 后，训练环境已经提供 NumPy、SciPy 和 MuJoCo。不同工具还需要： + +- 纯解析工具：Python、NumPy。 +- MuJoCo 扫描：`mujoco`、NumPy。 +- 完整 MPC：`pinocchio`、`osqp`、SciPy。 +- GUI：系统可用的 Tk/Tkinter。 + +Pinocchio 和 OSQP 没有加入训练环境锁文件，因为它们只服务于可选 MPC 工具，且 Pinocchio 的安装方式与操作系统、Conda/Python 环境有关。 + +## 常用命令 + +在 `05_software/train/rc_mjlab` 下执行： + +```bash +# 不启动 MuJoCo 的快速姿态表 +uv run python mujoco_sim/posture_tool.py +uv run python mujoco_sim/rl_friendly_opt.py + +# 姿态扫描 +uv run python mujoco_sim/posture_optimizer.py --analyze +uv run python mujoco_sim/static_posture_optimizer.py --quick + +# IK 与差速轮参数快速扫描 +uv run python mujoco_sim/ik_diff_sweep.py --quick + +# 完整 GUI/MPC 仿真，需要可选依赖 +uv run python mujoco_sim/run.py +``` + +## 参数边界 + +这是一份历史工具快照，保留当时用于分析和调参的常量： + +- `config.py`、`posture_optimizer.py` 和 `static_posture_optimizer.py` 中的解析质量常量为 `12.3 kg`。 +- 当前新版 MJCF 的惯性质量合计约为 `18.0377 kg`。 +- `config.py` 的姿态表默认值为髋俯仰 `0.666`、膝关节 `-1.546`；比赛训练架构的 Rough 默认姿态为 `0.550/-1.125`。 + +MuJoCo 直接加载模型的工具会使用 MJCF 内的质量和惯性；显式读取 `ROBOT_MASS` 的解析计算和 MPC 工具仍使用历史常量。使用输出作为新版本控制参数前，应先根据目标机械状态完成质量、惯性和默认姿态复核。 diff --git a/05_software/train/rc_mjlab/mujoco_sim/config.py b/05_software/train/rc_mjlab/mujoco_sim/config.py new file mode 100644 index 0000000..a314a85 --- /dev/null +++ b/05_software/train/rc_mjlab/mujoco_sim/config.py @@ -0,0 +1,81 @@ +"""Robot constants and control parameters.""" + +from pathlib import Path + +# Paths +REPO_ROOT = Path(__file__).resolve().parents[1] +SCENE_XML = REPO_ROOT / "mjcf" / "scene.xml" +MJCF_PATH = REPO_ROOT / "mjcf" / "wheelleg.xml" + +# Robot geometry +WHEEL_RADIUS = 0.10 # m +WHEEL_TRACK = 0.32 # m, left-right distance +ROBOT_MASS = 12.3 # kg +MAX_TORQUE = 17.0 # Nm per joint +MAX_JOINT_VEL = 13.0 # rad/s + +# Leg link lengths from MJCF, measured to wheel center. +L_THIGH = 0.25 +L_CALF = 0.20 + +# Leg names and joint ordering +LEG_NAMES = ("fl", "fr", "rl", "rr") +LEG_JOINTS = ("hip_abduction_joint", "hip_pitch_joint", "knee_joint") +WHEEL_JOINT = "wheel_joint" + +# Default standing pose aligned with the soft wheel-X height table. +# height ~= 0.37m, wheel x-offset ~= 0, peak/RMS leg torque balanced. +DEFAULT_JOINT_ANGLES = { + "hip_abduction": 0.0, + "hip_pitch": 0.666, + "knee": -1.546, +} + +# Actuator modes, configured at runtime: +# Leg joints: position PD, ctrl = target angle +# Wheel joints: velocity, ctrl = target velocity in rad/s + +# Control rates +SIM_DT = 0.002 +CTRL_DT = 0.02 +CTRL_DECIMATION = int(CTRL_DT / SIM_DT) + +# Wheel drive +WHEEL_VEL_MAX = 10.0 + +# Body pose control gains for height/roll/pitch compensation. +KP_HEIGHT = 3.0 +KP_ROLL = 0.5 +KP_PITCH = 0.5 + +# Calibrated height-to-joint-angle table. +# Constraint: avoid large wheel-center X offset from the hip/leg. This is a +# soft support-geometry guardrail, not a strict x=0 requirement. +# The optimizer also considers peak motor torque and RMS torque, so one hot +# motor is not hidden by a low average across all motors. +# Format: (height_m, hip_pitch_rad, knee_rad) +HEIGHT_TABLE = [ + (0.17, 0.914, -2.628), + (0.19, 0.926, -2.528), + (0.21, 0.924, -2.428), + (0.23, 0.912, -2.328), + (0.25, 0.892, -2.226), + (0.27, 0.864, -2.122), + (0.29, 0.834, -2.014), + (0.31, 0.798, -1.906), + (0.33, 0.758, -1.792), + (0.35, 0.714, -1.672), + (0.37, 0.666, -1.546), + (0.39, 0.612, -1.412), + (0.41, 0.552, -1.266), + (0.43, 0.484, -1.104), + (0.45, 0.404, -0.918), +] + +# Gait parameters +GAIT_FREQ = 2.5 +GAIT_DUTY = 0.6 +SWING_HEIGHT = 0.06 + +# Trot phase offsets: FL/RR in phase, FR/RL in phase +PHASE_OFFSETS = {"fl": 0.0, "fr": 0.5, "rl": 0.5, "rr": 0.0} diff --git a/05_software/train/rc_mjlab/mujoco_sim/controller.py b/05_software/train/rc_mjlab/mujoco_sim/controller.py new file mode 100644 index 0000000..396773a --- /dev/null +++ b/05_software/train/rc_mjlab/mujoco_sim/controller.py @@ -0,0 +1,367 @@ +"""Main controller: wheel mode + trot mode for wheeled-legged robot. + +Wheel mode: differential drive + leg posture hold (height/roll/pitch compensation) +Trot mode: quadruped gait with wheel-assisted propulsion + +Actuator interface: + - Leg joints: ctrl = target angle (PD: kp=60, kd=3) + - Wheel joints: ctrl = target velocity in rad/s (gain=2.0) +""" + +import numpy as np +from robot import Robot, RobotState +from dynamics import Dynamics +from mpc_controller import MPCController +from config import ( + LEG_NAMES, DEFAULT_JOINT_ANGLES, WHEEL_RADIUS, WHEEL_TRACK, + WHEEL_VEL_MAX, KP_ROLL, KP_PITCH, HEIGHT_TABLE, + GAIT_FREQ, GAIT_DUTY, SWING_HEIGHT, PHASE_OFFSETS, +) + +RL_ROUGH_Q = np.array([0.0, 0.550, -1.125], dtype=float) +LEG_STATE_IDX = np.array([0, 1, 2, 4, 5, 6, 8, 9, 10, 12, 13, 14], dtype=int) + + +class Controller: + """Wheeled-legged robot controller.""" + + def __init__(self, robot: Robot): + self.robot = robot + self.dynamics = Dynamics() + + # User commands + self.vel_x = 0.0 # m/s forward + self.vel_y = 0.0 # m/s lateral + self.yaw_rate = 0.0 # rad/s + self.height = 0.37 # m desired body height (wheel center under hip) + self.wheel_posture = "table" # "table" follows height slider; "rl" matches src/robot default + + # Mode: "wheel", "trot", or "mpc" + self.mode = "wheel" + + # Prone (lie down) state + self.prone = False + + # MPC controller + self._mpc_ctrl = MPCController(robot) + self._mpc_active = False # track torque mode state + + # Gait state + self._gait_phase = 0.0 + + # Smoothed commands for trot mode (avoid sudden jumps) + self._smooth_vx = 0.0 + self._smooth_vy = 0.0 + self._smooth_yaw = 0.0 + self._last_leg_targets = np.tile( + np.array( + [ + DEFAULT_JOINT_ANGLES["hip_abduction"], + DEFAULT_JOINT_ANGLES["hip_pitch"], + DEFAULT_JOINT_ANGLES["knee"], + ], + dtype=float, + ), + 4, + ) + + # Wheel-mode sensor feedback. + self.yaw_rate_kp = 0.45 + self.roll_comp_gain = KP_ROLL + self.pitch_comp_gain = KP_PITCH + self.encoder_posture_kp = 0.12 + self.encoder_posture_max = 0.025 + self.encoder_guard_start = 0.28 + self.encoder_guard_stop = 0.65 + self.imu_guard_start = np.deg2rad(12.0) + self.imu_guard_stop = np.deg2rad(28.0) + self.yaw_wheel_gain = 1.0 + self.max_yaw_wheel_speed = 4.0 + + # Default leg angles + self._default_q = np.array([ + DEFAULT_JOINT_ANGLES["hip_abduction"], + DEFAULT_JOINT_ANGLES["hip_pitch"], + DEFAULT_JOINT_ANGLES["knee"], + ]) + + # Swing leg memory + self._swing_start_foot = {leg: np.zeros(3) for leg in LEG_NAMES} + self._last_contact = {leg: True for leg in LEG_NAMES} + + def compute(self, state: RobotState, dt: float) -> tuple[np.ndarray, np.ndarray]: + # Smooth all velocity commands (both modes) + alpha = min(dt * 3.0, 1.0) # ~0.33s time constant + self._smooth_vx += alpha * (self.vel_x - self._smooth_vx) + self._smooth_vy += alpha * (self.vel_y - self._smooth_vy) + self._smooth_yaw += alpha * (self.yaw_rate - self._smooth_yaw) + + if self.prone: + self._ensure_position_mode() + return self._prone_mode() + if self.mode == "mpc": + return self._mpc_mode(state, dt) + if self.mode == "wheel": + self._ensure_position_mode() + return self._wheel_mode(state, dt) + else: + self._ensure_position_mode() + return self._trot_mode(state, dt) + + def _mpc_mode(self, state: RobotState, dt: float): + """MPC locomotion: MIT motor protocol (PD + MPC feedforward torque).""" + # Switch to torque mode if not already + if not self._mpc_active: + self.robot.enable_torque_mode() + self._mpc_active = True + + # Sync commands to MPC controller + self._mpc_ctrl.vel_x = self.vel_x + self._mpc_ctrl.vel_y = self.vel_y + self._mpc_ctrl.yaw_rate = self.yaw_rate + self._mpc_ctrl.height = self.height + + # Compute and apply (sets ctrl directly via set_ctrl_mit) + self._mpc_ctrl.compute(state, dt) + # Return dummy - ctrl already set + return np.zeros(12), np.zeros(4) + + def _ensure_position_mode(self): + """Switch back to position PD mode if coming from MPC.""" + if self._mpc_active: + self.robot.enable_position_mode() + self._mpc_active = False + + def _prone_mode(self): + """Lie down: actual prone pose from real robot.""" + leg_targets = np.zeros(12) + for i, leg in enumerate(LEG_NAMES): + side = 1.0 if leg[1] == "l" else -1.0 + leg_targets[i*3] = side * 0.3 # fl/rl: +0.3, fr/rr: -0.3 + leg_targets[i*3+1] = 1.5 # hip pitch + leg_targets[i*3+2] = -2.65 # knee hard limit from MJCF + return leg_targets, np.zeros(4) + + # ───────────────────────────────────────────────────────────────────── + # WHEEL MODE + # ───────────────────────────────────────────────────────────────────── + + def _wheel_mode(self, state: RobotState, dt: float): + """Wheel drive + leg posture hold. + + vel_y: limited effect in wheel mode (differential drive cannot produce + pure lateral motion). Uses hip_abduction lean for small lateral force. + For significant lateral motion, use trot mode. + """ + leg_targets = self._posture_control(state) + safe_vx, safe_yaw = self._wheel_velocity_envelope(self._smooth_vx, self._smooth_yaw) + yaw_feedback = safe_yaw + self.yaw_rate_kp * (safe_yaw - float(state.ang_vel[2])) + wheel_targets = self._differential_drive(safe_vx, yaw_feedback) + wheel_targets *= self._sensor_command_scale(state, leg_targets) + self._last_leg_targets = leg_targets.copy() + return leg_targets, wheel_targets + + def _posture_control(self, state: RobotState) -> np.ndarray: + """Leg joint targets from the soft wheel-X height table.""" + leg_targets = np.zeros(12) + + # Calibrated height→angle lookup (minimizes motor torque at each height) + _H = [r[0] for r in HEIGHT_TABLE] + _HIP = [r[1] for r in HEIGHT_TABLE] + _KNEE = [r[2] for r in HEIGHT_TABLE] + + h_clamp = np.clip(self.height, _H[0], _H[-1]) + q_hip_base = float(np.interp(h_clamp, _H, _HIP)) + q_knee_base = float(np.interp(h_clamp, _H, _KNEE)) + q_ab_base = 0.0 + if self.wheel_posture == "rl": + q_ab_base, q_hip_base, q_knee_base = RL_ROUGH_Q + + roll_corr = -self.roll_comp_gain * state.rpy[0] + pitch_corr = -self.pitch_comp_gain * state.rpy[1] + lateral_lean = 0.3 * self.vel_y + + for i, leg in enumerate(LEG_NAMES): + side = 1.0 if leg[1] == "l" else -1.0 + leg_targets[i*3] = np.clip(q_ab_base + side * roll_corr + lateral_lean, -0.5, 0.5) + leg_targets[i*3+1] = np.clip(q_hip_base + pitch_corr, -1.0, 2.5) + leg_targets[i*3+2] = np.clip(q_knee_base, -2.65, -0.3) + + encoder_err = self._last_leg_targets - state.joint_pos[LEG_STATE_IDX] + leg_targets += np.clip( + self.encoder_posture_kp * encoder_err, + -self.encoder_posture_max, + self.encoder_posture_max, + ) + leg_targets[0::3] = np.clip(leg_targets[0::3], -0.5, 0.5) + leg_targets[1::3] = np.clip(leg_targets[1::3], -1.0, 2.5) + leg_targets[2::3] = np.clip(leg_targets[2::3], -2.65, -0.3) + return leg_targets + + def _wheel_velocity_envelope(self, vel_x: float, yaw_rate: float) -> tuple[float, float]: + """Limit x/yaw combinations that are unsafe for the RL posture.""" + ax = abs(vel_x) + if ax >= 0.8: + yaw_lim = 0.35 + elif ax >= 0.5: + yaw_lim = 0.55 + elif ax >= 0.25: + yaw_lim = 0.75 + else: + yaw_lim = 1.0 + return float(vel_x), float(np.clip(yaw_rate, -yaw_lim, yaw_lim)) + + def _sensor_command_scale(self, state: RobotState, leg_targets: np.ndarray) -> float: + """Back off wheels when IMU or encoder feedback says posture is degrading.""" + leg_error = float(np.max(np.abs(state.joint_pos[LEG_STATE_IDX] - leg_targets))) + tilt = float(np.hypot(state.rpy[0], state.rpy[1])) + scale = 1.0 + + if leg_error >= self.encoder_guard_stop: + scale = 0.0 + elif leg_error > self.encoder_guard_start: + span = max(1e-6, self.encoder_guard_stop - self.encoder_guard_start) + scale *= 1.0 - (leg_error - self.encoder_guard_start) / span + + if tilt >= self.imu_guard_stop: + scale = 0.0 + elif tilt > self.imu_guard_start: + span = max(1e-6, self.imu_guard_stop - self.imu_guard_start) + scale *= 1.0 - (tilt - self.imu_guard_start) / span + + return float(np.clip(scale, 0.0, 1.0)) + + # ───────────────────────────────────────────────────────────────────── + # TROT MODE + # ───────────────────────────────────────────────────────────────────── + + def _trot_mode(self, state: RobotState, dt: float): + """Trot gait with wheel assist.""" + # Advance gait phase + self._gait_phase = (self._gait_phase + dt * GAIT_FREQ) % 1.0 + + # Contact state + contacts = {} + for leg in LEG_NAMES: + phase = (self._gait_phase + PHASE_OFFSETS[leg]) % 1.0 + contacts[leg] = phase < GAIT_DUTY + + # Pinocchio update + q_pin, dq_pin = self.robot.get_qpos_qvel_for_pinocchio() + self.dynamics.update(q_pin, dq_pin) + + leg_targets = np.zeros(12) + wheel_targets = np.zeros(4) + + for i, leg in enumerate(LEG_NAMES): + if contacts[leg]: + # Stance: posture hold + leg_targets[i*3:(i+1)*3] = self._stance_leg_target(state, leg) + self._swing_start_foot[leg] = self.dynamics.get_foot_pos(leg) + self._last_contact[leg] = True + # Wheel: drive with smoothed velocity + wheel_targets[i] = self._differential_drive_single( + self._smooth_vx, self._smooth_yaw, leg) + else: + # Swing: IK trajectory + swing_phase = self._get_swing_phase(leg) + target_foot = self._compute_swing_target(leg, state, swing_phase) + q_ik = self.dynamics.inverse_kinematics(leg, target_foot, q_pin) + leg_targets[i*3:(i+1)*3] = q_ik + self._last_contact[leg] = False + # Wheel: zero (free during swing) + wheel_targets[i] = 0.0 + + return leg_targets, wheel_targets + + def _stance_leg_target(self, state: RobotState, leg: str) -> np.ndarray: + """Stance leg: table-interpolated height + attitude compensation.""" + _H = [r[0] for r in HEIGHT_TABLE] + _HIP = [r[1] for r in HEIGHT_TABLE] + _KNEE = [r[2] for r in HEIGHT_TABLE] + + h_clamp = np.clip(self.height, _H[0], _H[-1]) + q_hip = float(np.interp(h_clamp, _H, _HIP)) + q_knee = float(np.interp(h_clamp, _H, _KNEE)) + + roll_corr = -KP_ROLL * state.rpy[0] + pitch_corr = -KP_PITCH * state.rpy[1] + side = 1.0 if leg[1] == "l" else -1.0 + lateral_lean = 0.3 * self.vel_y + + return np.array([ + np.clip(side * roll_corr + lateral_lean, -0.5, 0.5), + np.clip(q_hip + pitch_corr, -1.0, 2.5), + np.clip(q_knee, -2.65, -0.3), + ]) + + # ───────────────────────────────────────────────────────────────────── + # DIFFERENTIAL DRIVE + # ───────────────────────────────────────────────────────────────────── + + def _differential_drive(self, vel_x: float, yaw_rate: float) -> np.ndarray: + """4 wheel velocities from body commands.""" + linear_wheel = vel_x / WHEEL_RADIUS + yaw_wheel = self.yaw_wheel_gain * 0.5 * WHEEL_TRACK * yaw_rate / WHEEL_RADIUS + yaw_wheel = float(np.clip(yaw_wheel, -self.max_yaw_wheel_speed, self.max_yaw_wheel_speed)) + vel_left = linear_wheel - yaw_wheel + vel_right = linear_wheel + yaw_wheel + targets = np.zeros(4) + for i, leg in enumerate(LEG_NAMES): + targets[i] = vel_left if leg[1] == "l" else vel_right + return np.clip(targets, -WHEEL_VEL_MAX, WHEEL_VEL_MAX) + + def _differential_drive_single(self, vel_x: float, yaw_rate: float, leg: str) -> float: + if leg[1] == "l": + v = (vel_x - 0.5 * WHEEL_TRACK * yaw_rate) / WHEEL_RADIUS + else: + v = (vel_x + 0.5 * WHEEL_TRACK * yaw_rate) / WHEEL_RADIUS + return np.clip(v, -WHEEL_VEL_MAX, WHEEL_VEL_MAX) + + # ───────────────────────────────────────────────────────────────────── + # SWING TRAJECTORY + # ───────────────────────────────────────────────────────────────────── + + def _get_swing_phase(self, leg: str) -> float: + phase = (self._gait_phase + PHASE_OFFSETS[leg]) % 1.0 + if phase < GAIT_DUTY: + return 0.0 + return (phase - GAIT_DUTY) / (1.0 - GAIT_DUTY) + + def _compute_swing_target(self, leg: str, state: RobotState, + swing_phase: float) -> np.ndarray: + """Swing foot target with Raibert heuristic using COMMANDED velocity.""" + p_start = self._swing_start_foot[leg] + p_end = self._compute_touchdown(leg, state) + + s = swing_phase + s_mj = 10*s**3 - 15*s**4 + 6*s**5 + + pos = p_start + (p_end - p_start) * s_mj + + # Z lift + z_lift = 64.0 * s**3 * (1.0 - s)**3 + pos[2] = p_start[2] + SWING_HEIGHT * z_lift + + return pos + + def _compute_touchdown(self, leg: str, state: RobotState) -> np.ndarray: + """Raibert heuristic using COMMANDED velocity. + + When commands are zero, foot lands at its takeoff position (no net motion). + When commands are nonzero, foot placement is offset by commanded velocity. + """ + # Base: land where the foot took off (zero net displacement) + td = self._swing_start_foot[leg].copy() + + # Add commanded velocity offset (Raibert-style) + t_stance = (1.0 / GAIT_FREQ) * GAIT_DUTY + yaw = state.rpy[2] + c, s = np.cos(yaw), np.sin(yaw) + R_z = np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]]) + cmd_vel_world = R_z @ np.array([self._smooth_vx, self._smooth_vy, 0.0]) + + td[0] += cmd_vel_world[0] * t_stance * 0.5 + td[1] += cmd_vel_world[1] * t_stance * 0.5 + td[2] = WHEEL_RADIUS # ground level + return td diff --git a/05_software/train/rc_mjlab/mujoco_sim/dynamics.py b/05_software/train/rc_mjlab/mujoco_sim/dynamics.py new file mode 100644 index 0000000..d4f1e22 --- /dev/null +++ b/05_software/train/rc_mjlab/mujoco_sim/dynamics.py @@ -0,0 +1,97 @@ +"""Pinocchio dynamics: FK, Jacobian, IK for the wheeled-legged robot.""" + +import numpy as np +import pinocchio as pin +from config import MJCF_PATH, LEG_NAMES + +# Foot frame names in Pinocchio model (wheel link centers) +FOOT_FRAMES = {leg: f"{leg}_wheel_Link" for leg in LEG_NAMES} + +# Leg joint names for each leg +_LEG_JOINT_NAMES = { + leg: [f"{leg}_{jt}" for jt in ("hip_abduction_joint", "hip_pitch_joint", "knee_joint")] + for leg in LEG_NAMES +} + + +class Dynamics: + """Pinocchio-based kinematics/dynamics. Deployable on real hardware.""" + + def __init__(self): + self.model = pin.buildModelFromMJCF(str(MJCF_PATH)) + self.data = self.model.createData() + + # Cache frame IDs + self._foot_fids = {} + for leg, fname in FOOT_FRAMES.items(): + self._foot_fids[leg] = self.model.getFrameId(fname) + + # Cache joint velocity indices for each leg (3 joints) + self._leg_v_indices = {} + for leg, jnames in _LEG_JOINT_NAMES.items(): + indices = [] + for jn in jnames: + jid = self.model.getJointId(jn) + indices.append(self.model.joints[jid].idx_v) + self._leg_v_indices[leg] = indices + + # Cache joint config indices for each leg + self._leg_q_indices = {} + for leg, jnames in _LEG_JOINT_NAMES.items(): + indices = [] + for jn in jnames: + jid = self.model.getJointId(jn) + indices.append(self.model.joints[jid].idx_q) + self._leg_q_indices[leg] = indices + + def update(self, q: np.ndarray, dq: np.ndarray): + """Forward kinematics + Jacobians. + + Args: + q: Pinocchio config (nq=23: pos3, quat_xyzw4, joints16) + dq: Pinocchio velocity (nv=22: v_body3, w_body3, joints16) + """ + pin.forwardKinematics(self.model, self.data, q, dq) + pin.updateFramePlacements(self.model, self.data) + pin.computeJointJacobians(self.model, self.data, q) + + def get_foot_pos(self, leg: str) -> np.ndarray: + """Foot (wheel center) position in world frame (3,).""" + return self.data.oMf[self._foot_fids[leg]].translation.copy() + + def get_foot_jacobian_leg(self, leg: str) -> np.ndarray: + """3x3 linear Jacobian of foot w.r.t. 3 leg joints (world frame).""" + fid = self._foot_fids[leg] + J_full = pin.getFrameJacobian( + self.model, self.data, fid, pin.LOCAL_WORLD_ALIGNED)[:3, :] + cols = self._leg_v_indices[leg] + return J_full[:, cols] + + def inverse_kinematics(self, leg: str, target_pos: np.ndarray, + q_current: np.ndarray, max_iter=30, eps=1e-4) -> np.ndarray: + """Numerical IK for one leg. Returns (3,) joint angles. + + Args: + leg: Leg name + target_pos: Desired foot position in world frame (3,) + q_current: Current full Pinocchio config (nq=23) + """ + q = q_current.copy() + fid = self._foot_fids[leg] + q_indices = self._leg_q_indices[leg] + + for _ in range(max_iter): + pin.forwardKinematics(self.model, self.data, q) + pin.updateFramePlacements(self.model, self.data) + err = target_pos - self.data.oMf[fid].translation + if np.linalg.norm(err) < eps: + break + pin.computeJointJacobians(self.model, self.data, q) + J = pin.getFrameJacobian( + self.model, self.data, fid, pin.LOCAL_WORLD_ALIGNED)[:3, :] + J_leg = J[:, self._leg_v_indices[leg]] + dq = np.linalg.solve(J_leg.T @ J_leg + 1e-6 * np.eye(3), J_leg.T @ err) + for i, idx in enumerate(q_indices): + q[idx] += dq[i] + + return np.array([q[idx] for idx in q_indices]) diff --git a/05_software/train/rc_mjlab/mujoco_sim/gui.py b/05_software/train/rc_mjlab/mujoco_sim/gui.py new file mode 100644 index 0000000..c405aad --- /dev/null +++ b/05_software/train/rc_mjlab/mujoco_sim/gui.py @@ -0,0 +1,127 @@ +"""GUI control panel for the wheeled-legged robot.""" + +import tkinter as tk +from tkinter import ttk + + +class GUI: + """Tkinter control panel: sliders + gait buttons + status display.""" + + def __init__(self, controller): + self.ctrl = controller + self.root = tk.Tk() + self.root.title("WheelLeg Control") + self.root.geometry("400x500") + self.root.protocol("WM_DELETE_WINDOW", self._on_close) + self._closed = False + + self._build() + + def _build(self): + # Mode buttons + mf = ttk.LabelFrame(self.root, text="Mode") + mf.pack(fill="x", padx=8, pady=4) + for mode in ("wheel", "trot", "mpc"): + ttk.Button(mf, text=mode.upper(), + command=lambda m=mode: self._set_mode(m) + ).pack(side="left", padx=4, expand=True) + ttk.Button(mf, text="PRONE/STAND", + command=self._toggle_prone).pack(side="left", padx=4, expand=True) + + # Command sliders + cf = ttk.LabelFrame(self.root, text="Commands") + cf.pack(fill="x", padx=8, pady=4) + + self.vel_x_var = tk.DoubleVar(value=0.0) + self.vel_y_var = tk.DoubleVar(value=0.0) + self.yaw_var = tk.DoubleVar(value=0.0) + self.height_var = tk.DoubleVar(value=self.ctrl.height) + + self._slider(cf, "Vel X", self.vel_x_var, -1.5, 1.5) + self._slider(cf, "Vel Y*", self.vel_y_var, -0.5, 0.5) + self._slider(cf, "Yaw", self.yaw_var, -2.0, 2.0) + self._slider(cf, "Height", self.height_var, 0.17, 0.45) + + ttk.Label(cf, text="* Vel Y: trot mode only (diff-drive can't sidestep)", + font=("", 8)).pack(anchor="w", padx=8) + + ttk.Button(cf, text="Reset", command=self._reset).pack(pady=4) + + # Status display + sf = ttk.LabelFrame(self.root, text="Status") + sf.pack(fill="both", expand=True, padx=8, pady=4) + self.status_text = tk.Text(sf, height=12, width=45, font=("Consolas", 9)) + self.status_text.pack(fill="both", expand=True, padx=4, pady=4) + + def _slider(self, parent, label, var, lo, hi): + f = ttk.Frame(parent) + f.pack(fill="x", padx=4, pady=2) + ttk.Label(f, text=label, width=7).pack(side="left") + ttk.Scale(f, from_=lo, to=hi, variable=var, + command=lambda *_: self._sync()).pack(side="left", fill="x", expand=True) + lbl = ttk.Label(f, text="0.00", width=6) + lbl.pack(side="left") + var.trace_add("write", lambda *_, v=var, l=lbl: l.config(text=f"{v.get():.2f}")) + + def _set_mode(self, mode): + self.ctrl.mode = mode + self.ctrl.prone = False + + def _toggle_prone(self): + self.ctrl.prone = not self.ctrl.prone + + def _sync(self): + self.ctrl.vel_x = self.vel_x_var.get() + self.ctrl.vel_y = self.vel_y_var.get() + self.ctrl.yaw_rate = self.yaw_var.get() + self.ctrl.height = self.height_var.get() + + def _reset(self): + self.vel_x_var.set(0.0) + self.vel_y_var.set(0.0) + self.yaw_var.set(0.0) + self._sync() + + def _on_close(self): + self._closed = True + self.root.destroy() + + @property + def closed(self): + return self._closed + + def update_status(self, state, step): + """Update status text with current robot state.""" + txt = ( + f"Mode: {self.ctrl.mode} Step: {step}\n" + f"Pos: x={state.pos[0]:.3f} y={state.pos[1]:.3f} z={state.pos[2]:.3f}\n" + f"RPY: r={np.degrees(state.rpy[0]):.1f}° p={np.degrees(state.rpy[1]):.1f}° " + f"y={np.degrees(state.rpy[2]):.1f}°\n" + f"Vel: vx={state.lin_vel[0]:.3f} vy={state.lin_vel[1]:.3f} vz={state.lin_vel[2]:.3f}\n" + f"Cmd: vx={self.ctrl.vel_x:.2f} yaw={self.ctrl.yaw_rate:.2f} h={self.ctrl.height:.3f}\n" + f"─────────────────────────────────\n" + ) + # Joint angles (compact) + for i, leg in enumerate(("FL", "FR", "RL", "RR")): + q = state.joint_pos[i*4:i*4+3] + w = state.joint_vel[i*4+3] + txt += f"{leg}: [{q[0]:+.2f} {q[1]:+.2f} {q[2]:+.2f}] w={w:+.1f}\n" + + self.status_text.delete("1.0", tk.END) + self.status_text.insert(tk.END, txt) + + def tick(self): + """Process GUI events. Returns False if window closed.""" + if self._closed: + return False + try: + self.root.update_idletasks() + self.root.update() + return True + except tk.TclError: + self._closed = True + return False + + +# Need numpy for degrees conversion in update_status +import numpy as np diff --git a/05_software/train/rc_mjlab/mujoco_sim/ik_diff_sweep.py b/05_software/train/rc_mjlab/mujoco_sim/ik_diff_sweep.py new file mode 100644 index 0000000..22d941b --- /dev/null +++ b/05_software/train/rc_mjlab/mujoco_sim/ik_diff_sweep.py @@ -0,0 +1,389 @@ +"""Sweep wheel-mode IK postures for differential-drive tracking. + +The sweep is intentionally small and reproducible: + 1. Generate ab=0 leg postures in the requested height range. + 2. Keep candidates with good static geometry from rl_friendly_opt. + 3. Simulate forward, yaw, and arc commands in MuJoCo. + 4. Rank by attitude, x-speed tracking, yaw-rate tracking, and wheel contact. + +Usage: + uv run python mujoco_sim/ik_diff_sweep.py --quick + uv run python mujoco_sim/ik_diff_sweep.py --height-min 0.15 --height-max 0.42 +""" + +from __future__ import annotations + +import argparse +import json +import math +import sys +from dataclasses import dataclass +from pathlib import Path + +import mujoco +import numpy as np + +THIS_DIR = Path(__file__).resolve().parent +if str(THIS_DIR) not in sys.path: + sys.path.insert(0, str(THIS_DIR)) + +from config import LEG_JOINTS, LEG_NAMES, SCENE_XML, WHEEL_JOINT, WHEEL_RADIUS # noqa: E402 +from rl_friendly_opt import get_all, rl_cost # noqa: E402 +from robot import Robot # noqa: E402 + + +@dataclass(frozen=True) +class Candidate: + height: float + ab: float + hip: float + knee: float + static_cost: float + r_hip_x: float + cond: float + max_tau: float + + +@dataclass(frozen=True) +class Trial: + name: str + vx: float + yaw_rate: float + duration: float + + +def wrap_pi(x: float) -> float: + return (x + math.pi) % (2.0 * math.pi) - math.pi + + +def body_track(model: mujoco.MjModel, data: mujoco.MjData) -> float: + wheel_bids = [] + for leg in LEG_NAMES: + bid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, f"{leg}_wheel_Link") + wheel_bids.append(bid) + y = data.xipos[wheel_bids, 1] + return float(np.mean(y[[0, 2]]) - np.mean(y[[1, 3]])) + + +def build_maps(robot: Robot) -> tuple[dict[str, int], dict[str, int], dict[str, int]]: + act: dict[str, int] = {} + qadr: dict[str, int] = {} + vadr: dict[str, int] = {} + for leg in LEG_NAMES: + for jt in (*LEG_JOINTS, WHEEL_JOINT): + name = f"{leg}_{jt}" + act[name] = mujoco.mj_name2id(robot.model, mujoco.mjtObj.mjOBJ_ACTUATOR, name) + jid = mujoco.mj_name2id(robot.model, mujoco.mjtObj.mjOBJ_JOINT, name) + qadr[name] = robot.model.jnt_qposadr[jid] + vadr[name] = robot.model.jnt_dofadr[jid] + return act, qadr, vadr + + +def set_posture(robot: Robot, cand: Candidate, act: dict[str, int], qadr: dict[str, int]) -> None: + mujoco.mj_resetData(robot.model, robot.data) + for leg in LEG_NAMES: + side_ab = cand.ab if leg[1] == "l" else -cand.ab + for jt, val in zip(LEG_JOINTS, (side_ab, cand.hip, cand.knee)): + name = f"{leg}_{jt}" + robot.data.qpos[qadr[name]] = val + robot.data.ctrl[act[name]] = val + robot.data.ctrl[act[f"{leg}_{WHEEL_JOINT}"]] = 0.0 + robot.data.qpos[:3] = [0.0, 0.0, max(0.25, cand.height + 0.08)] + robot.data.qpos[3:7] = [1.0, 0.0, 0.0, 0.0] + robot.data.qvel[:] = 0.0 + mujoco.mj_forward(robot.model, robot.data) + + +def wheel_targets( + vx: float, + yaw_rate: float, + track: float, + max_wheel: float, + yaw_gain: float, + wheel_model: str, + linear_gain: float, + direct_yaw_gain: float, + wheel_signs: np.ndarray, +) -> np.ndarray: + if wheel_model == "direct": + left = linear_gain * vx - direct_yaw_gain * yaw_rate + right = linear_gain * vx + direct_yaw_gain * yaw_rate + else: + left = (vx - yaw_gain * 0.5 * track * yaw_rate) / WHEEL_RADIUS + right = (vx + yaw_gain * 0.5 * track * yaw_rate) / WHEEL_RADIUS + raw = np.array([left, right, left, right], dtype=float) + return np.clip(raw * wheel_signs, -max_wheel, max_wheel) + + +def run_trial(robot: Robot, cand: Candidate, trial: Trial, args: argparse.Namespace) -> dict: + act, qadr, vadr = build_maps(robot) + set_posture(robot, cand, act, qadr) + + ctrl_dt = args.control_dt + sim_dt = robot.model.opt.timestep + steps_per_ctrl = max(1, int(round(ctrl_dt / sim_dt))) + track = body_track(robot.model, robot.data) if args.track_source == "model" else args.track_width + wheel_signs = np.array(args.wheel_signs, dtype=float) + wheel_cmd = wheel_targets( + trial.vx, + trial.yaw_rate, + track, + args.max_wheel_speed, + args.yaw_gain, + args.wheel_model, + args.linear_gain, + args.direct_yaw_gain, + wheel_signs, + ) + + for _ in range(int(round(args.settle / sim_dt))): + for leg in LEG_NAMES: + vals = (cand.ab if leg[1] == "l" else -cand.ab, cand.hip, cand.knee) + for jt, val in zip(LEG_JOINTS, vals): + robot.data.ctrl[act[f"{leg}_{jt}"]] = val + robot.data.ctrl[act[f"{leg}_{WHEEL_JOINT}"]] = 0.0 + robot.step() + + state0 = robot.get_state() + yaw0 = float(state0.rpy[2]) + x0 = float(state0.pos[0]) + + max_roll = 0.0 + max_pitch = 0.0 + max_tilt = 0.0 + max_wheel_air = -1e9 + wheel_err_sum = 0.0 + samples = 0 + max_leg_err = 0.0 + body_vx_sum = 0.0 + yaw_unwrapped = 0.0 + last_yaw = yaw0 + leg_target = np.array([cand.ab, cand.hip, cand.knee] * 4, dtype=float) + + total_steps = int(round(trial.duration / sim_dt)) + cmd = np.zeros(4, dtype=float) + max_delta = args.wheel_accel_limit * ctrl_dt + wheel_body_ids = [ + mujoco.mj_name2id(robot.model, mujoco.mjtObj.mjOBJ_BODY, f"{leg}_wheel_Link") + for leg in LEG_NAMES + ] + wheel_vadr = np.array([vadr[f"{leg}_{WHEEL_JOINT}"] for leg in LEG_NAMES], dtype=int) + leg_qadr = np.array([qadr[f"{leg}_{jt}"] for leg in LEG_NAMES for jt in LEG_JOINTS], dtype=int) + + for step in range(total_steps): + if step % steps_per_ctrl == 0: + cmd = cmd + np.clip(wheel_cmd - cmd, -max_delta, max_delta) + for leg in LEG_NAMES: + vals = (cand.ab if leg[1] == "l" else -cand.ab, cand.hip, cand.knee) + for jt, val in zip(LEG_JOINTS, vals): + robot.data.ctrl[act[f"{leg}_{jt}"]] = val + for i, leg in enumerate(LEG_NAMES): + robot.data.ctrl[act[f"{leg}_{WHEEL_JOINT}"]] = cmd[i] + + robot.step() + + if step % steps_per_ctrl == 0: + st = robot.get_state() + roll, pitch = float(st.rpy[0]), float(st.rpy[1]) + yaw_now = float(st.rpy[2]) + yaw_unwrapped += wrap_pi(yaw_now - last_yaw) + last_yaw = yaw_now + body_vx_sum += float(st.rot[:, 0].dot(st.lin_vel)) + max_roll = max(max_roll, abs(roll)) + max_pitch = max(max_pitch, abs(pitch)) + max_tilt = max(max_tilt, math.hypot(roll, pitch)) + wheel_air = robot.data.xipos[wheel_body_ids, 2] - WHEEL_RADIUS + max_wheel_air = max(max_wheel_air, float(np.max(wheel_air))) + wheel_err_sum += float(np.mean(np.abs(robot.data.qvel[wheel_vadr] - cmd))) + max_leg_err = max(max_leg_err, float(np.max(np.abs(robot.data.qpos[leg_qadr] - leg_target)))) + samples += 1 + + st = robot.get_state() + elapsed = max(1e-6, float(st.time - state0.time)) + world_x_rate = (float(st.pos[0]) - x0) / elapsed + x_rate = body_vx_sum / max(1, samples) + yaw_rate = yaw_unwrapped / max(1e-6, samples * steps_per_ctrl * sim_dt) + x_err = abs(x_rate - trial.vx) + yaw_err = abs(yaw_rate - trial.yaw_rate) + return { + "trial": trial.name, + "x_rate": x_rate, + "world_x_rate": world_x_rate, + "yaw_rate": yaw_rate, + "x_err": x_err, + "yaw_err": yaw_err, + "max_roll_deg": math.degrees(max_roll), + "max_pitch_deg": math.degrees(max_pitch), + "max_tilt_deg": math.degrees(max_tilt), + "max_wheel_air_m": max_wheel_air, + "mean_wheel_err": wheel_err_sum / max(1, samples), + "max_leg_err": max_leg_err, + "track": track, + "wheel_cmd": [float(x) for x in wheel_cmd], + } + + +def generate_candidates(args: argparse.Namespace) -> list[Candidate]: + if args.fixed_hip is not None or args.fixed_knee is not None: + if args.fixed_hip is None or args.fixed_knee is None: + raise SystemExit("--fixed-hip and --fixed-knee must be provided together") + ab = float(args.fixed_ab) + hip = float(args.fixed_hip) + knee = float(args.fixed_knee) + r = get_all(ab, hip, knee) + return [ + Candidate( + height=float(r["z"]), + ab=ab, + hip=hip, + knee=knee, + static_cost=float(rl_cost(r)), + r_hip_x=float(r["r_hip_x_mag"]), + cond=float(r["cond"]), + max_tau=float(r["max_tau"]), + ) + ] + + cands: list[Candidate] = [] + h_targets = np.arange(args.height_min, args.height_max + 0.5 * args.height_step, args.height_step) + ab_values = [0.0] if args.ab_max <= 1e-9 else np.arange(0.0, args.ab_max + 1e-9, args.ab_step) + hip_values = np.arange(args.hip_min, args.hip_max + 0.5 * args.hip_step, args.hip_step) + knee_values = np.arange(args.knee_min, args.knee_max + 0.5 * args.knee_step, args.knee_step) + for ht in h_targets: + bucket: list[Candidate] = [] + for ab in ab_values: + for hip in hip_values: + for knee in knee_values: + r = get_all(float(ab), float(hip), float(knee)) + if float(r["z"]) < args.height_min or float(r["z"]) > args.height_max: + continue + if abs(float(r["z"]) - float(ht)) > args.height_tol: + continue + if r["wz"] >= r["kz"]: + continue + if abs(r["r_hip_x_mag"]) > args.max_wheel_x: + continue + cost = float(rl_cost(r)) + bucket.append( + Candidate( + height=float(r["z"]), + ab=float(ab), + hip=float(hip), + knee=float(knee), + static_cost=cost, + r_hip_x=float(r["r_hip_x_mag"]), + cond=float(r["cond"]), + max_tau=float(r["max_tau"]), + ) + ) + bucket.sort(key=lambda c: c.static_cost) + cands.extend(bucket[: args.top_per_height]) + return cands + + +def score_result(cand: Candidate, trials: list[dict]) -> float: + score = 0.08 * cand.static_cost + for t in trials: + score += 8.0 * t["x_err"] + score += 10.0 * t["yaw_err"] + score += 0.08 * t["max_tilt_deg"] + score += 0.03 * max(0.0, t["max_pitch_deg"] - 8.0) ** 2 + score += 20.0 * max(0.0, t["max_wheel_air_m"] - 0.015) + score += 1.5 * t["mean_wheel_err"] + score += 2.0 * max(0.0, t["max_leg_err"] - 0.35) + return float(score) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--height-min", type=float, default=0.15) + parser.add_argument("--height-max", type=float, default=0.42) + parser.add_argument("--height-step", type=float, default=0.01) + parser.add_argument("--height-tol", type=float, default=0.004) + parser.add_argument("--top-per-height", type=int, default=1) + parser.add_argument("--ab-max", type=float, default=0.0) + parser.add_argument("--ab-step", type=float, default=0.04) + parser.add_argument("--hip-min", type=float, default=0.25) + parser.add_argument("--hip-max", type=float, default=1.05) + parser.add_argument("--hip-step", type=float, default=0.025) + parser.add_argument("--knee-min", type=float, default=-2.65) + parser.add_argument("--knee-max", type=float, default=-0.85) + parser.add_argument("--knee-step", type=float, default=0.025) + parser.add_argument("--max-wheel-x", type=float, default=0.09) + parser.add_argument("--fixed-ab", type=float, default=0.0) + parser.add_argument("--fixed-hip", type=float, default=None) + parser.add_argument("--fixed-knee", type=float, default=None) + parser.add_argument("--duration", type=float, default=4.0) + parser.add_argument("--settle", type=float, default=1.5) + parser.add_argument("--control-dt", type=float, default=0.02) + parser.add_argument("--vx", type=float, default=0.6) + parser.add_argument("--yaw", type=float, default=0.3) + parser.add_argument("--arc-yaw", type=float, default=0.15) + parser.add_argument("--yaw-gain", type=float, default=1.0) + parser.add_argument("--wheel-model", choices=("diff", "direct"), default="diff") + parser.add_argument("--linear-gain", type=float, default=12.5) + parser.add_argument("--direct-yaw-gain", type=float, default=8.0) + parser.add_argument("--max-wheel-speed", type=float, default=12.0) + parser.add_argument("--wheel-accel-limit", type=float, default=35.0) + parser.add_argument("--track-source", choices=("model", "fixed"), default="model") + parser.add_argument("--track-width", type=float, default=0.394) + parser.add_argument( + "--wheel-signs", + type=float, + nargs=4, + default=[1.0, 1.0, 1.0, 1.0], + metavar=("FL", "FR", "RL", "RR"), + help="Per-wheel velocity sign multipliers in joint order.", + ) + parser.add_argument("--quick", action="store_true") + parser.add_argument("--json", type=Path, default=None) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + if args.quick: + args.height_step = 0.02 + args.hip_step = 0.05 + args.knee_step = 0.05 + + candidates = generate_candidates(args) + if not candidates: + raise SystemExit("No candidates found") + + trials = [ + Trial("forward", args.vx, 0.0, args.duration), + Trial("yaw", 0.0, args.yaw, args.duration), + Trial("arc", args.vx, args.arc_yaw, args.duration), + ] + robot = Robot(SCENE_XML) + rows = [] + for i, cand in enumerate(candidates, 1): + trial_rows = [run_trial(robot, cand, t, args) for t in trials] + rows.append({"candidate": cand.__dict__, "trials": trial_rows, "score": score_result(cand, trial_rows)}) + if i % 10 == 0: + print(f"tested {i}/{len(candidates)}") + + rows.sort(key=lambda r: r["score"]) + if args.json: + args.json.write_text(json.dumps(rows, indent=2), encoding="utf-8") + + print("Top IK postures for differential drive tracking") + print("rank score height ab hip knee static xhip cond tau | fwd_x yaw_wz arc_x arc_wz max_tilt max_pitch") + for rank, row in enumerate(rows[:10], 1): + c = row["candidate"] + by = {t["trial"]: t for t in row["trials"]} + max_tilt = max(t["max_tilt_deg"] for t in row["trials"]) + max_pitch = max(t["max_pitch_deg"] for t in row["trials"]) + print( + f"{rank:>2} {row['score']:>7.2f} {c['height']:.3f} {c['ab']:.2f} {c['hip']:.3f} {c['knee']:.3f} " + f"{c['static_cost']:.1f} {c['r_hip_x']:.3f} {c['cond']:.2f} {c['max_tau']:.2f} | " + f"{by['forward']['x_rate']:.3f} {by['yaw']['yaw_rate']:.3f} {by['arc']['x_rate']:.3f} {by['arc']['yaw_rate']:.3f} " + f"{max_tilt:.1f} {max_pitch:.1f}" + ) + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/05_software/train/rc_mjlab/mujoco_sim/mpc.py b/05_software/train/rc_mjlab/mujoco_sim/mpc.py new file mode 100644 index 0000000..3491651 --- /dev/null +++ b/05_software/train/rc_mjlab/mujoco_sim/mpc.py @@ -0,0 +1,242 @@ +"""Convex MPC solver for wheeled-legged robot. + +Centroidal dynamics: single rigid body model with 4 contact forces. +State: x = [pos(3), rpy(3), vel(3), omega(3)] = 12 +Input: u = [f1(3), f2(3), f3(3), f4(3)] = 12 +Friction pyramid constraints on each foot. + +Reference: MIT Cheetah 3 Convex MPC (Di Carlo et al.) +""" + +import numpy as np +from scipy import sparse +from scipy.linalg import block_diag +import osqp + +from config import ROBOT_MASS, LEG_NAMES + +# MPC parameters +MPC_HORIZON = 10 # prediction steps +MPC_DT = 0.02 # 50 Hz MPC update +MU = 0.6 # friction coefficient +FZ_MAX = 200.0 # max vertical force per leg +FZ_MIN = 10.0 # min vertical force (stance) +NX = 12 # state dim +NU = 12 # input dim (4 legs × 3D force) + +# Cost weights: [pos_x, pos_y, pos_z, roll, pitch, yaw, vx, vy, vz, wx, wy, wz] +Q_WEIGHTS = np.array([2.0, 2.0, 50.0, 50.0, 50.0, 10.0, 2.0, 2.0, 1.0, 1.0, 1.0, 1.0]) +R_WEIGHTS = np.array([1e-6] * 12) + + +def _skew(v): + return np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]]) + + +class ConvexMPC: + """Convex MPC: solves QP for optimal ground reaction forces.""" + + def __init__(self, mass=ROBOT_MASS, inertia=None): + self.mass = mass + # Approximate body inertia (diagonal, world-aligned) + if inertia is None: + self.I_body = np.diag([0.07, 0.26, 0.24]) + else: + self.I_body = np.array(inertia).reshape(3, 3) + + self.N = MPC_HORIZON + self.dt = MPC_DT + self.Q = np.diag(Q_WEIGHTS) + self.R = np.diag(R_WEIGHTS) + self.gravity = np.array([0, 0, -9.81]) + + self._last_forces = np.zeros(NU) + + def solve(self, x0, x_ref, foot_positions, contact_schedule): + """Solve MPC QP. + + Args: + x0: (12,) current state [pos, rpy, vel, omega] + x_ref: (12, N) reference trajectory over horizon + foot_positions: (4, 3) foot positions in world frame (relative to CoM) + contact_schedule: (4, N) binary contact table (1=stance) + + Returns: + forces: (12,) optimal forces for current timestep [f1x,f1y,f1z,...,f4x,f4y,f4z] + """ + N = self.N + + # Build dynamics matrices + Ad, Bd_list, gd = self._discretize_dynamics(x0, foot_positions) + + # Build QP: min 0.5 z'Hz + f'z s.t. lb <= Az <= ub, lbx <= z <= ubx + # Decision variables: z = [x1,...,xN, u0,...,uN-1] + nvars = N * NX + N * NU + + # --- Hessian --- + H_diag = np.concatenate([np.tile(2 * Q_WEIGHTS, N), np.tile(2 * R_WEIGHTS, N)]) + H = sparse.diags(H_diag, format='csc') + + # --- Gradient --- + g = np.zeros(nvars) + for k in range(N): + g[k*NX:(k+1)*NX] = -2 * self.Q @ x_ref[:, k] + + # --- Dynamics equality constraints --- + # x_{k+1} = Ad @ x_k + Bd_k @ u_k + gd + # Rewrite: x_{k+1} - Ad @ x_k - Bd_k @ u_k = gd (for k>0) + # x_1 - Bd_0 @ u_0 = Ad @ x0 + gd (for k=0) + n_eq = N * NX + A_eq = np.zeros((n_eq, nvars)) + b_eq = np.zeros(n_eq) + + # k=0: x_1 = Ad @ x0 + Bd_0 @ u_0 + gd + A_eq[0:NX, 0:NX] = np.eye(NX) # x_1 + A_eq[0:NX, N*NX:N*NX+NU] = -Bd_list[0] # -Bd_0 @ u_0 + b_eq[0:NX] = Ad @ x0 + gd + + for k in range(1, N): + row = k * NX + # x_{k+1} + A_eq[row:row+NX, k*NX:(k+1)*NX] = np.eye(NX) + # -Ad @ x_k + A_eq[row:row+NX, (k-1)*NX:k*NX] = -Ad + # -Bd_k @ u_k + A_eq[row:row+NX, N*NX+k*NU:N*NX+(k+1)*NU] = -Bd_list[k] + b_eq[row:row+NX] = gd + + # --- Friction pyramid inequality constraints --- + # For each stance leg at each timestep: 4 faces + # fx - mu*fz <= 0, -fx - mu*fz <= 0, fy - mu*fz <= 0, -fy - mu*fz <= 0 + n_ineq_max = 4 * 4 * N + A_ineq = np.zeros((n_ineq_max, nvars)) + u_ineq = np.zeros(n_ineq_max) + + row = 0 + for k in range(N): + u_base = N * NX + k * NU + for leg in range(4): + if contact_schedule[leg, k] == 1: + fx_idx = u_base + leg * 3 + fy_idx = u_base + leg * 3 + 1 + fz_idx = u_base + leg * 3 + 2 + + # Friction pyramid: stance leg + A_ineq[row, fx_idx] = 1.0 + A_ineq[row, fz_idx] = -MU + row += 1 + A_ineq[row, fx_idx] = -1.0 + A_ineq[row, fz_idx] = -MU + row += 1 + A_ineq[row, fy_idx] = 1.0 + A_ineq[row, fz_idx] = -MU + row += 1 + A_ineq[row, fy_idx] = -1.0 + A_ineq[row, fz_idx] = -MU + row += 1 + + A_ineq = A_ineq[:row] + u_ineq = u_ineq[:row] + + # Stack constraints + A_full = np.vstack([A_eq, A_ineq]) + l_full = np.concatenate([b_eq, -np.inf * np.ones(row)]) + u_full = np.concatenate([b_eq, u_ineq]) + + # --- Box constraints on forces (as identity rows in A) --- + A_box = np.zeros((N * NU, nvars)) + l_box = -np.inf * np.ones(N * NU) + u_box = np.inf * np.ones(N * NU) + + for k in range(N): + u_base = N * NX + k * NU + for leg in range(4): + idx = u_base + leg * 3 + box_row = k * NU + leg * 3 + # Identity rows for fx, fy, fz + for j in range(3): + A_box[box_row + j, idx + j] = 1.0 + + if contact_schedule[leg, k] == 1: + # Stance: fz bounded + l_box[box_row + 2] = FZ_MIN + u_box[box_row + 2] = FZ_MAX + else: + # Swing: all forces = 0 + l_box[box_row:box_row+3] = 0.0 + u_box[box_row:box_row+3] = 0.0 + + # Final constraint matrix + A_full = np.vstack([A_full, A_box]) + l_full = np.concatenate([l_full, l_box]) + u_full = np.concatenate([u_full, u_box]) + + # --- Solve with OSQP --- + A_sparse = sparse.csc_matrix(A_full) + H_sparse = sparse.triu(H, format='csc') + + solver = osqp.OSQP() + solver.setup(H_sparse, g, A_sparse, l_full, u_full, + eps_abs=1e-4, eps_rel=1e-4, + max_iter=500, polish=True, verbose=False, + warm_start=True) + + # Warm start with previous solution + if self._last_forces is not None: + x_warm = np.zeros(nvars) + x_warm[N*NX:N*NX+NU] = self._last_forces + solver.warm_start(x=x_warm) + + result = solver.solve() + + if result.info.status == 'solved' or result.info.status == 'solved_inaccurate': + # Extract first timestep forces + forces = result.x[N*NX:N*NX+NU] + self._last_forces = forces.copy() + else: + forces = self._last_forces + + return forces + + def _discretize_dynamics(self, x0, foot_positions): + """Build discrete-time centroidal dynamics. + + State: [pos, rpy, vel, omega] (12) + Continuous: dx/dt = Ac @ x + Bc @ u + gc + Discrete: x_{k+1} = Ad @ x + Bd @ u + gd + """ + m = self.mass + I_inv = np.linalg.inv(self.I_body) + dt = self.dt + yaw = x0[5] + cy, sy = np.cos(yaw), np.sin(yaw) + + # Rotation for rpy rate ≈ R_z^T @ omega + R_zT = np.array([[cy, sy, 0], [-sy, cy, 0], [0, 0, 1]]) + + # Ac (12×12) + Ac = np.zeros((NX, NX)) + Ac[0:3, 6:9] = np.eye(3) # pos_dot = vel + Ac[3:6, 9:12] = R_zT # rpy_dot ≈ R_z^T @ omega + + # Ad = I + Ac*dt (first-order) + Ad = np.eye(NX) + Ac * dt + + # Bc varies per timestep (foot positions change contact point) + Bd_list = [] + for k in range(self.N): + Bc = np.zeros((NX, NU)) + for leg in range(4): + r = foot_positions[leg] + # vel_dot += f/m + Bc[6:9, leg*3:(leg+1)*3] = np.eye(3) / m + # omega_dot += I^{-1} @ (r × f) + Bc[9:12, leg*3:(leg+1)*3] = I_inv @ _skew(r) + Bd = Bc * dt + Bd_list.append(Bd) + + # Gravity contribution + gd = np.zeros(NX) + gd[6:9] = self.gravity * dt # vel += g*dt + + return Ad, Bd_list, gd diff --git a/05_software/train/rc_mjlab/mujoco_sim/mpc_controller.py b/05_software/train/rc_mjlab/mujoco_sim/mpc_controller.py new file mode 100644 index 0000000..0cbff00 --- /dev/null +++ b/05_software/train/rc_mjlab/mujoco_sim/mpc_controller.py @@ -0,0 +1,261 @@ +"""MPC controller integration for wheeled-legged robot. + +Integrates: gait scheduler + reference trajectory + ConvexMPC solver + + swing leg control + stance force mapping + wheel drive. + +Architecture (following go2-convex-mpc): + - MPC runs at ~50 Hz (every MPC_DECIMATION control steps) + - Swing/stance leg controller runs at control rate (50 Hz) + - Wheel drive: stance legs use differential drive, swing legs coast +""" + +import numpy as np +from robot import Robot, RobotState +from dynamics import Dynamics +from mpc import ConvexMPC, MPC_DT +from config import ( + LEG_NAMES, WHEEL_RADIUS, WHEEL_TRACK, WHEEL_VEL_MAX, + CTRL_DT, GAIT_FREQ, GAIT_DUTY, SWING_HEIGHT, PHASE_OFFSETS, + DEFAULT_JOINT_ANGLES, ROBOT_MASS, +) + +# MPC update decimation (relative to control loop) +MPC_DECIMATION = max(1, int(MPC_DT / CTRL_DT)) # 1 step at 50Hz + + +class MPCController: + """Convex MPC locomotion controller for wheeled-legged robot.""" + + def __init__(self, robot: Robot): + self.robot = robot + self.dynamics = Dynamics() + self.mpc = ConvexMPC(mass=ROBOT_MASS) + + # User commands + self.vel_x = 0.0 + self.vel_y = 0.0 + self.yaw_rate = 0.0 + self.height = 0.35 # actual standing height with default joint angles + + # Gait state - start at phase 0 with all legs in stance (duty=0.6) + self._gait_phase = 0.0 + self._step_count = 0 + self._initialized = False + + # MPC solution cache - initialize with gravity compensation + self._mpc_forces = np.zeros(12) + self._init_gravity_comp() + + # Swing trajectory state + self._swing_start_foot = {leg: np.zeros(3) for leg in LEG_NAMES} + self._swing_start_time = {leg: 0.0 for leg in LEG_NAMES} + self._last_contact = {leg: True for leg in LEG_NAMES} + + # Smoothed commands + self._smooth_vx = 0.0 + self._smooth_vy = 0.0 + self._smooth_yaw = 0.0 + + def _init_gravity_comp(self): + """Pre-fill MPC forces with static gravity compensation.""" + fz_per_leg = ROBOT_MASS * 9.81 / 4.0 + for i in range(4): + self._mpc_forces[i*3 + 2] = fz_per_leg + + def compute(self, state: RobotState, dt: float): + """Main MPC control loop. + + Uses MIT motor protocol: tau = kp*(q_des-q) + kd*(dq_des-dq) + tau_ff + where tau_ff comes from MPC force mapping via Jacobian transpose. + + Returns: + tau_legs: (12,) feedforward torques for MIT mode + wheel_targets: (4,) wheel velocity targets + """ + # Smooth commands + alpha = min(dt * 3.0, 1.0) + self._smooth_vx += alpha * (self.vel_x - self._smooth_vx) + self._smooth_vy += alpha * (self.vel_y - self._smooth_vy) + self._smooth_yaw += alpha * (self.yaw_rate - self._smooth_yaw) + + # Update Pinocchio + q_pin, dq_pin = self.robot.get_qpos_qvel_for_pinocchio() + self.dynamics.update(q_pin, dq_pin) + + # Initialize foot positions on first call + if not self._initialized: + for leg in LEG_NAMES: + self._swing_start_foot[leg] = self.dynamics.get_foot_pos(leg) + self._initialized = True + + # Decide if we should trot or just stand + moving = (abs(self._smooth_vx) > 0.02 or + abs(self._smooth_vy) > 0.02 or + abs(self._smooth_yaw) > 0.05) + + if moving: + self._gait_phase = (self._gait_phase + dt * GAIT_FREQ) % 1.0 + else: + self._gait_phase = 0.0 # all legs in stance + + # Contact schedule + contacts = {} + for leg in LEG_NAMES: + phase = (self._gait_phase + PHASE_OFFSETS[leg]) % 1.0 + contacts[leg] = phase < GAIT_DUTY + + # Get foot positions relative to CoM + foot_positions = np.zeros((4, 3)) + for i, leg in enumerate(LEG_NAMES): + foot_positions[i] = self.dynamics.get_foot_pos(leg) - state.pos + + # --- Run MPC at lower rate --- + if self._step_count % MPC_DECIMATION == 0: + x0 = self._build_state_vector(state) + x_ref = self._build_reference(state) + contact_table = self._build_contact_table() + self._mpc_forces = self.mpc.solve(x0, x_ref, foot_positions, contact_table) + + self._step_count += 1 + + # --- Compute feedforward torques and desired joint positions --- + tau_ff = np.zeros(12) + q_des = np.zeros(12) + dq_des = np.zeros(12) + kp = np.zeros(12) + kd = np.zeros(12) + wheel_targets = np.zeros(4) + + for i, leg in enumerate(LEG_NAMES): + if contacts[leg]: + # Stance: MPC force → feedforward torque, PD holds posture + f_leg = self._mpc_forces[i*3:(i+1)*3] + J = self.dynamics.get_foot_jacobian_leg(leg) + tau_ff[i*3:(i+1)*3] = J.T @ (-f_leg) + + # PD target: default standing angles (posture hold) + q_des[i*3] = DEFAULT_JOINT_ANGLES["hip_abduction"] + q_des[i*3+1] = DEFAULT_JOINT_ANGLES["hip_pitch"] + q_des[i*3+2] = DEFAULT_JOINT_ANGLES["knee"] + kp[i*3:(i+1)*3] = [40.0, 40.0, 40.0] + kd[i*3:(i+1)*3] = [3.0, 3.0, 3.0] + + # Record foot position + self._swing_start_foot[leg] = self.dynamics.get_foot_pos(leg) + self._last_contact[leg] = True + + # Wheel drive + wheel_targets[i] = self._wheel_cmd(leg) + else: + # Swing: IK target position, strong PD, no feedforward + if self._last_contact[leg]: + self._swing_start_foot[leg] = self.dynamics.get_foot_pos(leg) + self._swing_start_time[leg] = state.time + self._last_contact[leg] = False + + q_ik = self._swing_leg_ik(leg, state, q_pin) + q_des[i*3:(i+1)*3] = q_ik + kp[i*3:(i+1)*3] = [60.0, 60.0, 60.0] # strong PD for swing + kd[i*3:(i+1)*3] = [3.0, 3.0, 3.0] + # tau_ff stays 0 for swing + + wheel_targets[i] = 0.0 + + # Use MIT protocol via robot interface + self.robot.set_ctrl_mit(q_des, dq_des, kp, kd, tau_ff, wheel_targets) + # Return dummy (actual ctrl is set directly above) + return None, None + + def _build_state_vector(self, state: RobotState): + """Build MPC state: [pos, rpy, vel, omega].""" + return np.concatenate([state.pos, state.rpy, state.lin_vel, state.ang_vel]) + + def _build_reference(self, state: RobotState): + """Build reference trajectory over MPC horizon.""" + N = self.mpc.N + x_ref = np.zeros((12, N)) + + yaw = state.rpy[2] + cy, sy = np.cos(yaw), np.sin(yaw) + R_z = np.array([[cy, -sy, 0], [sy, cy, 0], [0, 0, 1]]) + vel_world = R_z @ np.array([self._smooth_vx, self._smooth_vy, 0.0]) + + for k in range(N): + t = (k + 1) * self.mpc.dt + # Position: integrate from current + x_ref[0, k] = state.pos[0] + vel_world[0] * t + x_ref[1, k] = state.pos[1] + vel_world[1] * t + x_ref[2, k] = self.height + # RPY: keep roll/pitch zero, integrate yaw + x_ref[3, k] = 0.0 + x_ref[4, k] = 0.0 + x_ref[5, k] = yaw + self._smooth_yaw * t + # Velocity + x_ref[6, k] = vel_world[0] + x_ref[7, k] = vel_world[1] + x_ref[8, k] = 0.0 + # Angular velocity + x_ref[9, k] = 0.0 + x_ref[10, k] = 0.0 + x_ref[11, k] = self._smooth_yaw + + return x_ref + + def _build_contact_table(self): + """Build contact schedule over MPC horizon.""" + N = self.mpc.N + table = np.zeros((4, N), dtype=int) + for k in range(N): + future_phase = (self._gait_phase + (k + 1) * self.mpc.dt * GAIT_FREQ) % 1.0 + for i, leg in enumerate(LEG_NAMES): + leg_phase = (future_phase + PHASE_OFFSETS[leg]) % 1.0 + table[i, k] = 1 if leg_phase < GAIT_DUTY else 0 + return table + + def _swing_leg_ik(self, leg: str, state: RobotState, q_pin: np.ndarray): + """Swing leg: compute IK target joint angles for trajectory.""" + swing_phase = self._get_swing_phase(leg) + + p_start = self._swing_start_foot[leg] + p_end = self._compute_touchdown(leg, state) + + s = swing_phase + s_mj = 10*s**3 - 15*s**4 + 6*s**5 + + pos_des = p_start + (p_end - p_start) * s_mj + # Z lift + z_lift = 64.0 * s**3 * (1.0 - s)**3 + pos_des[2] = p_start[2] + SWING_HEIGHT * z_lift + + # IK to get joint angles + q_ik = self.dynamics.inverse_kinematics(leg, pos_des, q_pin) + return q_ik + + def _get_swing_phase(self, leg: str) -> float: + phase = (self._gait_phase + PHASE_OFFSETS[leg]) % 1.0 + if phase < GAIT_DUTY: + return 0.0 + return (phase - GAIT_DUTY) / (1.0 - GAIT_DUTY) + + def _compute_touchdown(self, leg: str, state: RobotState) -> np.ndarray: + """Raibert heuristic for touchdown position.""" + td = self._swing_start_foot[leg].copy() + t_stance = GAIT_DUTY / GAIT_FREQ + + yaw = state.rpy[2] + cy, sy = np.cos(yaw), np.sin(yaw) + R_z = np.array([[cy, -sy, 0], [sy, cy, 0], [0, 0, 1]]) + cmd_vel_world = R_z @ np.array([self._smooth_vx, self._smooth_vy, 0.0]) + + td[0] += cmd_vel_world[0] * t_stance * 0.5 + td[1] += cmd_vel_world[1] * t_stance * 0.5 + td[2] = WHEEL_RADIUS + return td + + def _wheel_cmd(self, leg: str) -> float: + """Differential drive for a single wheel.""" + if leg[1] == "l": + v = (self._smooth_vx - 0.5 * WHEEL_TRACK * self._smooth_yaw) / WHEEL_RADIUS + else: + v = (self._smooth_vx + 0.5 * WHEEL_TRACK * self._smooth_yaw) / WHEEL_RADIUS + return np.clip(v, -WHEEL_VEL_MAX, WHEEL_VEL_MAX) diff --git a/05_software/train/rc_mjlab/mujoco_sim/posture_optimizer.py b/05_software/train/rc_mjlab/mujoco_sim/posture_optimizer.py new file mode 100644 index 0000000..a49565a --- /dev/null +++ b/05_software/train/rc_mjlab/mujoco_sim/posture_optimizer.py @@ -0,0 +1,367 @@ +"""Posture optimizer for wheeled-leg standing/crawl height table. + +The table is not a pure "minimum average torque" table. For crawl and low-bar +traversal, the wheel center should not be far from the hip/leg in sagittal X, +otherwise the robot is no longer really using the wheel as the support/drive +point. This is a soft guardrail, not a strict x=0 constraint. The score combines: + +1. wheel center X offset from hip +2. peak single-motor holding torque +3. RMS torque, used as a proxy for I^2R heating + +Usage: + python posture_optimizer.py # MuJoCo sweep + python posture_optimizer.py --quick # coarse MuJoCo sweep + python posture_optimizer.py --analyze # analytical-only sweep +""" + +import argparse +import math +import time +from pathlib import Path + +import numpy as np + +REPO_ROOT = Path(__file__).resolve().parents[1] + +try: + import mujoco +except ImportError: + mujoco = None + +SCENE_XML = REPO_ROOT / "mjcf" / "scene.xml" +WHEEL_RADIUS = 0.10 +HIP_Z_OFFSET = 0.054 +L1, L2 = 0.25, 0.20 +ROBOT_MASS = 12.3 +G = 9.81 +F_PER_LEG = ROBOT_MASS * G / 4.0 +MAX_TORQUE = 17.0 + +LEG_NAMES = ("fl", "fr", "rl", "rr") +LEG_JOINTS = ("hip_abduction_joint", "hip_pitch_joint", "knee_joint") + +# 0.15m is not a good default table target for ab=0.0. 0.17m is the practical +# default crawl height, while lower crawl can be evaluated with abduction. +KNEE_MIN = -2.65 +HEIGHT_MIN = 0.17 +HEIGHT_MAX = 0.46 +SOFT_WHEEL_X_OFFSET = 0.05 +HARD_WHEEL_X_OFFSET = 0.08 + +_ACTUATOR_NAMES = [f"{leg}_{jt}" for leg in LEG_NAMES for jt in LEG_JOINTS] + + +def compute_fk(hip, knee): + """Return wheel-center x offset and base height from (hip_pitch, knee).""" + x = L1 * math.sin(hip) + L2 * math.sin(hip + knee) + z = L1 * math.cos(hip) + L2 * math.cos(hip + knee) + base_height = WHEEL_RADIUS + z - HIP_Z_OFFSET + return x, base_height + + +def posture_cost(x_foot, torques): + """Score one posture by support geometry, peak torque, and RMS torque.""" + tau = np.asarray(torques, dtype=float) + peak_torque = float(np.max(np.abs(tau))) + rms_torque = float(np.sqrt(np.mean(np.square(tau)))) + mean_i2r = float(np.mean(np.square(tau))) + x_penalty = max(0.0, abs(x_foot) - SOFT_WHEEL_X_OFFSET) + cost = ( + 0.5 * (abs(x_foot) / SOFT_WHEEL_X_OFFSET) ** 2 + + 8.0 * (x_penalty / max(1e-6, HARD_WHEEL_X_OFFSET - SOFT_WHEEL_X_OFFSET)) ** 2 + + 3.0 * (peak_torque / MAX_TORQUE) ** 2 + + (rms_torque / MAX_TORQUE) ** 2 + ) + return cost, peak_torque, rms_torque, mean_i2r + + +def analyze_analytical(): + """Analytical sweep using static GRF moments.""" + hip_range = np.arange(0.0, 1.6, 0.002) + knee_range = np.arange(KNEE_MIN, -0.4, 0.002) + + results = [] + for hip in hip_range: + for knee in knee_range: + x_foot, height = compute_fk(hip, knee) + if height < HEIGHT_MIN or height > HEIGHT_MAX: + continue + + tau_hip = F_PER_LEG * x_foot + x_knee_to_foot = L2 * math.sin(hip + knee) + tau_knee = F_PER_LEG * x_knee_to_foot + tau_abduction = 0.0 + cost, peak, rms, mean_i2r = posture_cost( + x_foot, (tau_abduction, tau_hip, tau_knee)) + + results.append({ + "hip": float(hip), + "knee": float(knee), + "height": float(height), + "x_foot": float(x_foot), + "tau_abd": tau_abduction, + "tau_hip": float(tau_hip), + "tau_knee": float(tau_knee), + "tau_peak": peak, + "tau_rms": rms, + "mean_i2r": mean_i2r, + "cost": cost, + }) + + return results + + +def run_mujoco_sweep(quick=False): + """MuJoCo sweep measuring actual actuator forces at steady state.""" + if mujoco is None: + raise RuntimeError("mujoco is not installed; use --analyze for analytical mode") + + model = mujoco.MjModel.from_xml_path(str(SCENE_XML)) + data = mujoco.MjData(model) + + act_ids = {} + for name in _ACTUATOR_NAMES: + act_ids[name] = mujoco.mj_name2id( + model, mujoco.mjtObj.mjOBJ_ACTUATOR, name) + + for i in range(model.nu): + name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_ACTUATOR, i) + model.actuator_biastype[i] = 1 + model.actuator_gaintype[i] = 0 + model.actuator_forcelimited[i] = 0 + if "wheel" not in name: + model.actuator_gainprm[i, 0] = 60.0 + model.actuator_biasprm[i, 0] = 0.0 + model.actuator_biasprm[i, 1] = -60.0 + model.actuator_biasprm[i, 2] = -3.0 + model.actuator_ctrlrange[i] = [-3.14, 3.14] + else: + model.actuator_gainprm[i, 0] = 2.0 + model.actuator_biasprm[i, 0] = 0.0 + model.actuator_biasprm[i, 1] = 0.0 + model.actuator_biasprm[i, 2] = -2.0 + model.actuator_ctrlrange[i] = [-20.0, 20.0] + + jnt_ids = {} + for name in _ACTUATOR_NAMES: + jid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, name) + jnt_ids[name] = model.jnt_qposadr[jid] + + if quick: + hip_range = np.arange(0.3, 1.5, 0.10) + knee_range = np.arange(KNEE_MIN, -0.6, 0.10) + else: + hip_range = np.arange(0.0, 1.6, 0.04) + knee_range = np.arange(KNEE_MIN, -0.4, 0.04) + + results = [] + total = 0 + valid = 0 + + for hip in hip_range: + for knee in knee_range: + total += 1 + x_foot, height = compute_fk(hip, knee) + if height < HEIGHT_MIN or height > HEIGHT_MAX: + continue + + mujoco.mj_resetData(model, data) + + for leg in LEG_NAMES: + for jt, val in zip(LEG_JOINTS, (0.0, hip, knee)): + name = f"{leg}_{jt}" + data.qpos[jnt_ids[name]] = val + + data.qpos[2] = height + data.qpos[3] = 1.0 + data.qpos[4:7] = 0.0 + + for leg in LEG_NAMES: + for jt, val in zip(LEG_JOINTS, (0.0, hip, knee)): + name = f"{leg}_{jt}" + data.ctrl[act_ids[name]] = val + name_w = f"{leg}_wheel_joint" + if name_w in act_ids: + data.ctrl[act_ids[name_w]] = 0.0 + + mujoco.mj_forward(model, data) + + for _ in range(500): + mujoco.mj_step(model, data) + + roll, pitch, _ = _get_rpy(data, model) + if abs(roll) > 0.8 or abs(pitch) > 0.8: + continue + + torque_buf = [] + for _ in range(100): + mujoco.mj_step(model, data) + torque_buf.append([data.actuator_force[act_ids[name]] + for name in _ACTUATOR_NAMES]) + tau_avg = np.array(torque_buf).mean(axis=0) + + tau_hip_val = tau_avg[1] + tau_knee_val = tau_avg[2] + tau_abd_val = tau_avg[0] + cost, peak, rms, mean_i2r = posture_cost(x_foot, tau_avg) + + q_hip_actual = float(data.qpos[jnt_ids["fl_hip_pitch_joint"]]) + q_knee_actual = float(data.qpos[jnt_ids["fl_knee_joint"]]) + + valid += 1 + results.append({ + "hip": float(hip), + "knee": float(knee), + "height": float(f"{height:.4f}"), + "x_foot": float(f"{x_foot:.4f}"), + "tau_abd": float(f"{tau_abd_val:.4f}"), + "tau_hip": float(f"{tau_hip_val:.4f}"), + "tau_knee": float(f"{tau_knee_val:.4f}"), + "tau_peak": float(f"{peak:.4f}"), + "tau_rms": float(f"{rms:.4f}"), + "mean_i2r": float(f"{mean_i2r:.4f}"), + "cost": float(f"{cost:.4f}"), + "q_hip_actual": float(f"{q_hip_actual:.4f}"), + "q_knee_actual": float(f"{q_knee_actual:.4f}"), + }) + + if valid % 20 == 0: + print(f" [{valid}/{total}] hip={hip:.2f} knee={knee:.2f} " + f"h={height:.3f} x={x_foot:+.4f} " + f"peak={peak:.2f} rms={rms:.2f} cost={cost:.4f}") + + return results + + +def _get_rpy(data, model): + """Extract roll and pitch from MuJoCo data.""" + base_bid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "base_link") + rot = data.xmat[base_bid].reshape(3, 3) + roll = math.atan2(rot[2, 1], rot[2, 2]) + pitch = math.atan2(-rot[2, 0], math.sqrt(rot[2, 1] ** 2 + rot[2, 2] ** 2)) + return roll, pitch, 0.0 + + +def print_top_results(results, n=10): + """Print the top-N results with lowest cost.""" + sorted_r = sorted(results, key=lambda r: r["cost"]) + + print(f"\n{'=' * 104}") + print(f"TOP {n} CONFIGURATIONS (soft wheel-X + peak/RMS torque score)") + print(f"{'=' * 104}") + print(f"{'Rank':>4} {'hip':>6} {'knee':>7} {'height':>7} {'x_foot':>8} " + f"{'tau_abd':>8} {'tau_hip':>8} {'tau_knee':>8} " + f"{'peak':>8} {'rms':>8} {'cost':>9}") + print(f"{'-' * 104}") + + for i, r in enumerate(sorted_r[:n]): + print(f"{i + 1:>4} {r['hip']:>6.3f} {r['knee']:>7.3f} " + f"{r['height']:>7.3f} {r.get('x_foot', 0):>8.4f} " + f"{r.get('tau_abd', 0):>8.3f} {r['tau_hip']:>8.3f} " + f"{r['tau_knee']:>8.3f} {r.get('tau_peak', 0):>8.3f} " + f"{r.get('tau_rms', 0):>8.3f} {r['cost']:>9.4f}") + + best = sorted_r[0] + print(f"\nBEST: hip={best['hip']:.3f} knee={best['knee']:.3f} " + f"z={best['height']:.3f}m x={best.get('x_foot', 0):+.4f}m " + f"peak={best.get('tau_peak', 0):.3f}Nm " + f"rms={best.get('tau_rms', 0):.3f}Nm cost={best['cost']:.4f}\n") + + return sorted_r + + +def compute_height_table(results): + """Build height-to-angle lookup with soft wheel-X support guardrail.""" + sorted_r = sorted( + (r for r in results if abs(r.get("x_foot", 999.0)) <= HARD_WHEEL_X_OFFSET), + key=lambda r: r["height"], + ) + if not sorted_r: + raise RuntimeError("No candidates satisfy HARD_WHEEL_X_OFFSET") + + h_range = np.arange(0.17, 0.46, 0.02) + table_h, table_hip, table_knee = [], [], [] + + for h_target in h_range: + candidates = [(r, abs(r["height"] - h_target)) for r in sorted_r] + candidates.sort(key=lambda x: (x[1], x[0]["cost"])) + best = candidates[0][0] + table_h.append(best["height"]) + table_hip.append(best["hip"]) + table_knee.append(best["knee"]) + + return { + "height": [round(h, 3) for h in table_h], + "hip": [round(h, 3) for h in table_hip], + "knee": [round(k, 3) for k in table_knee], + } + + +def export_calibrated_table(table): + """Print the new height table in copy-paste format.""" + print(f"\n{'=' * 80}") + print("CALIBRATED HEIGHT TABLE (soft wheel-X support guardrail)") + print(f"{'=' * 80}") + print(f"_H = {table['height']}") + print(f"_HIP = {table['hip']}") + print(f"_KNEE = {table['knee']}") + print(f"{'=' * 80}\n") + + +def main(): + parser = argparse.ArgumentParser(description="Find wheeled-leg posture table") + parser.add_argument("--quick", action="store_true", help="Coarse MuJoCo sweep") + parser.add_argument("--analyze", action="store_true", help="Analytical only") + parser.add_argument("--mujoco", action="store_true", default=True, + help="Run MuJoCo simulation when available") + args = parser.parse_args() + + print("=" * 72) + print("WHEELED-LEG POSTURE OPTIMIZER") + print("=" * 72) + print(f"Robot mass: {ROBOT_MASS} kg, F_per_leg: {F_PER_LEG:.1f} N") + print(f"L1={L1}m, L2={L2}m, wheel_r={WHEEL_RADIUS}m") + print(f"Height range: [{HEIGHT_MIN}, {HEIGHT_MAX}] m") + print(f"Knee min hard limit: {KNEE_MIN} rad") + print(f"Soft wheel X offset: {SOFT_WHEEL_X_OFFSET} m") + print(f"Hard wheel X offset: {HARD_WHEEL_X_OFFSET} m\n") + + t0 = time.time() + if args.analyze or mujoco is None: + print("[Analytical mode]") + results = analyze_analytical() + else: + print("[MuJoCo simulation mode]") + results = run_mujoco_sweep(quick=args.quick) + + elapsed = time.time() - t0 + print(f"Evaluated {len(results)} valid configurations in {elapsed:.1f}s") + if not results: + print("No valid configurations found") + return + + best_results = print_top_results(results, n=15) + + x_def, h_def = compute_fk(0.666, -1.546) + print(f"Current config default: hip=0.666, knee=-1.546 " + f"=> z={h_def:.3f}m, x={x_def:+.4f}m") + + table = compute_height_table(best_results) + export_calibrated_table(table) + + best = best_results[0] + print("=" * 72) + print("RECOMMENDED DEFAULT") + print("=" * 72) + print(f"hip_abduction: 0.0") + print(f"hip_pitch: {best['hip']:.3f}") + print(f"knee: {best['knee']:.3f}") + print(f"height: {best['height']:.3f} m") + print(f"x_foot: {best.get('x_foot', 0):+.4f} m") + print(f"tau_peak: {best.get('tau_peak', 0):.3f} Nm") + print(f"tau_rms: {best.get('tau_rms', 0):.3f} Nm") + + +if __name__ == "__main__": + main() diff --git a/05_software/train/rc_mjlab/mujoco_sim/posture_tool.py b/05_software/train/rc_mjlab/mujoco_sim/posture_tool.py new file mode 100644 index 0000000..f5428c9 --- /dev/null +++ b/05_software/train/rc_mjlab/mujoco_sim/posture_tool.py @@ -0,0 +1,189 @@ +"""Fast posture table helper based on wheelleg.xml link offsets. + +This tool is useful because it uses the real FL leg offsets from MJCF instead +of the simplified two-link geometry used by posture_optimizer.py. It is still a +static single-leg approximation, so use it to choose candidate crawl/standing +poses, then verify in MuJoCo and on the robot at low speed. +""" + +import math + +F_PER_LEG = 12.3 * 9.81 / 4.0 +KNEE_MIN = -2.65 +HIP_MIN = -2.58 +HIP_MAX = 2.58 +AB_MIN = -0.436 +AB_MAX = 0.611 +HARD_WHEEL_X_OFFSET = 0.08 + + +def rx(a): + c, s = math.cos(a), math.sin(a) + return ((1, 0, 0), (0, c, -s), (0, s, c)) + + +def ry(a): + c, s = math.cos(a), math.sin(a) + return ((c, 0, s), (0, 1, 0), (-s, 0, c)) + + +def mv(m, v): + return [ + m[0][0] * v[0] + m[0][1] * v[1] + m[0][2] * v[2], + m[1][0] * v[0] + m[1][1] * v[1] + m[1][2] * v[2], + m[2][0] * v[0] + m[2][1] * v[1] + m[2][2] * v[2], + ] + + +def add(a, b): + return [a[0] + b[0], a[1] + b[1], a[2] + b[2]] + + +def get_posture(q_ab, q_hip, q_knee): + """Return static FK/torque metrics for one FL leg. + + Returns: + base_z, tau_ab, tau_hip, tau_knee, i2r_total, knee_z, wheel_z, wheel_x_from_hip + """ + # FL offsets from mjcf/wheelleg.xml. + t_ab = [0.32826, 0.066172, 0.053981] + t_hip = [0.06389, -0.027344, 0.00010727] + t_knee = [0.0, 0.1035, -0.25] + t_wheel = [0.0, 0.014699, -0.20011] + + p = mv(ry(q_knee), t_wheel) + p = add(p, t_knee) + p = mv(ry(q_hip), p) + p = add(p, t_hip) + p = mv(rx(q_ab), p) + p = add(p, t_ab) + + base_z = 0.10 - p[2] + wheel_z = p[2] + + knee_pos = mv(ry(q_hip), t_knee) + knee_pos = mv(rx(q_ab), knee_pos) + knee_z = knee_pos[2] + t_ab[2] + + joint_knee = mv(ry(q_hip), add(t_knee, t_hip)) + joint_knee = mv(rx(q_ab), joint_knee) + joint_knee = add(t_ab, joint_knee) + + joint_hip = mv(rx(q_ab), t_hip) + joint_hip = add(t_ab, joint_hip) + + r_knee = [p[0] - joint_knee[0], p[1] - joint_knee[1], p[2] - joint_knee[2]] + r_hip = [p[0] - joint_hip[0], p[1] - joint_hip[1], p[2] - joint_hip[2]] + r_ab = [p[0] - t_ab[0], p[1] - t_ab[1], p[2] - t_ab[2]] + + tau_ab = r_ab[1] * F_PER_LEG + tau_hip = -r_hip[0] * F_PER_LEG + tau_knee = -r_knee[0] * F_PER_LEG + i2r_total = tau_ab * tau_ab + tau_hip * tau_hip + tau_knee * tau_knee + wheel_x_from_hip = p[0] - joint_hip[0] + return base_z, tau_ab, tau_hip, tau_knee, i2r_total, knee_z, wheel_z, wheel_x_from_hip + + +def score_candidate(z, z_target, tau_ab, tau_hip, tau_knee, wheel_x, x_target=0.0): + peak = max(abs(tau_ab), abs(tau_hip), abs(tau_knee)) + rms = math.sqrt((tau_ab * tau_ab + tau_hip * tau_hip + tau_knee * tau_knee) / 3.0) + x_err = wheel_x - x_target + x_over = max(0.0, abs(wheel_x) - 0.05) + return ( + 3000.0 * (z - z_target) ** 2 + + 2.5 * (peak / 17.0) ** 2 + + (rms / 17.0) ** 2 + + 0.4 * (x_err / 0.05) ** 2 + + 6.0 * (x_over / 0.03) ** 2 + ), peak, rms + + +def find_best(z_target, ab_range=(0.0, 0.0), step=0.002, x_target=0.0, hard_wheel_x_offset=HARD_WHEEL_X_OFFSET): + """Find one static posture near target height without violating hard limits.""" + best = None + ab0, ab1 = ab_range + n_ab = max(1, int(round((ab1 - ab0) / step)) + 1) + n_hip = int(round((1.6 - 0.0) / step)) + 1 + n_knee = int(round((-0.4 - KNEE_MIN) / step)) + 1 + + for ia in range(n_ab): + ab = ab0 + ia * step + if ab < AB_MIN or ab > AB_MAX: + continue + for ih in range(n_hip): + hip = ih * step + if hip < HIP_MIN or hip > HIP_MAX: + continue + for ik in range(n_knee): + knee = KNEE_MIN + ik * step + z, ta, th, tk, i2r, kz, wz, wx = get_posture(ab, hip, knee) + if abs(z - z_target) > 0.0015: + continue + if abs(wx) > hard_wheel_x_offset: + continue + if wz >= kz: + continue + cost, peak, rms = score_candidate(z, z_target, ta, th, tk, wx, x_target=x_target) + cand = (cost, ab, hip, knee, z, ta, th, tk, peak, rms, i2r, wx) + if best is None or cand[0] < best[0]: + best = cand + return best + + +def print_table( + z_targets, + name, + ab_range=(0.0, 0.0), + step=0.002, + x_target=0.0, + hard_wheel_x_offset=HARD_WHEEL_X_OFFSET, +): + print(f"\n{'=' * 96}") + print(name) + print(f"{'=' * 96}") + print(f"{'z':>6} {'ab':>6} {'hip':>7} {'knee':>7} " + f"{'tau_ab':>8} {'tau_hip':>8} {'tau_knee':>9} " + f"{'peak':>8} {'rms':>8} {'x_hip':>8}") + print("-" * 96) + for zt in z_targets: + best = find_best( + zt, + ab_range=ab_range, + step=step, + x_target=x_target, + hard_wheel_x_offset=hard_wheel_x_offset, + ) + if best is None: + print(f"{zt:>6.3f} no valid config") + continue + _, ab, hip, knee, z, ta, th, tk, peak, rms, _, wx = best + print(f"{z:>6.3f} {ab:>6.3f} {hip:>7.3f} {knee:>7.3f} " + f"{ta:>8.3f} {th:>8.3f} {tk:>9.3f} " + f"{peak:>8.3f} {rms:>8.3f} {wx:>8.4f}") + + +def print_crawl_default(): + best = find_best(0.17, ab_range=(0.0, 0.0), step=0.002) + if best is None: + return + _, ab, hip, knee, z, ta, th, tk, peak, rms, _, wx = best + print("\nSuggested runtime crawl_default_dof_pos:") + print( + f"[{ab:.3f}, {hip:.3f}, {knee:.3f}, " + f"{ab:.3f}, {hip:.3f}, {knee:.3f}, " + f"{ab:.3f}, {hip:.3f}, {knee:.3f}, " + f"{ab:.3f}, {hip:.3f}, {knee:.3f}, " + "0.0, 0.0, 0.0, 0.0]" + ) + print(f"# z={z:.3f}, peak={peak:.3f}Nm, rms={rms:.3f}Nm, wheel_x_from_hip={wx:+.4f}m") + + +if __name__ == "__main__": + print_table([round(0.36 + 0.01 * i, 2) for i in range(10)], + "STANDING candidates from MJCF FL geometry", step=0.004) + print_table([round(0.17 + 0.01 * i, 2) for i in range(9)], + "CRAWL candidates from MJCF FL geometry", step=0.002) + print_table([round(0.10 + 0.01 * i, 2) for i in range(6)], + "LOW CRAWL candidates, fixed FL abduction = +0.2", + ab_range=(0.2, 0.2), step=0.002, hard_wheel_x_offset=0.45) + print_crawl_default() diff --git a/05_software/train/rc_mjlab/mujoco_sim/rl_friendly_opt.py b/05_software/train/rc_mjlab/mujoco_sim/rl_friendly_opt.py new file mode 100644 index 0000000..32cad32 --- /dev/null +++ b/05_software/train/rc_mjlab/mujoco_sim/rl_friendly_opt.py @@ -0,0 +1,200 @@ +"""RlFriendlyPostureOpt — 结合电机发热 + RL友好度约束的站/爬姿优化 + +RL友好度约束（从实机经验总结，静力学可计算）： + 1. 髋关节力臂 ≥ 0.08m — 不让髋闲置（动态响应差） + 2. 三电机不均衡 ≤ 1.3x — 不让单电机先超载 + 3. 运动学条件数 κ ≤ 3.5 — 不让有效传动比过高（放大控制噪声） + +使用方法： + uv run python mujoco_sim/rl_friendly_opt.py # 打印推荐 + uv run python -c "from mujoco_sim.rl_friendly_opt import get_all; print(get_all(0, 0.8, -1.22))" + +基于 mjcf/wheelleg.xml 的 FL 腿运动学。 +""" +import math +import numpy as np + +# --------------------------------------------------------------------------- +# 运动学常数（来自 wheelleg.xml FL 腿） +# --------------------------------------------------------------------------- +T_AB = [0.32826, 0.066172, 0.053981] +T_HIP = [0.06389, -0.027344, 0.00010727] +T_KNEE = [0.0, 0.1035, -0.25] +T_WHEEL = [0.0, 0.014699, -0.20011] +F_PER_LEG = 12.3 * 9.81 / 4.0 +KNEE_MIN, KNEE_MAX = -2.65, 2.65 +HIP_MIN, HIP_MAX = -2.58, 2.58 + +# --------------------------------------------------------------------------- +# 工具函数 +# --------------------------------------------------------------------------- +def rx(a): + c = math.cos(a); s = math.sin(a) + return ((1, 0, 0), (0, c, -s), (0, s, c)) + +def ry(a): + c = math.cos(a); s = math.sin(a) + return ((c, 0, s), (0, 1, 0), (-s, 0, c)) + +def mv(m, v): + return [m[0][0]*v[0] + m[0][1]*v[1] + m[0][2]*v[2], + m[1][0]*v[0] + m[1][1]*v[1] + m[1][2]*v[2], + m[2][0]*v[0] + m[2][1]*v[1] + m[2][2]*v[2]] + +def add(a, b): return [a[0]+b[0], a[1]+b[1], a[2]+b[2]] + +# --------------------------------------------------------------------------- +# 主函数 +# --------------------------------------------------------------------------- +def get_all(q_ab, q_hip, q_knee): + """FK + 力矩 + 运动学指标。 + + 返回 dict: + z, ab, hip, knee, tau_ab, tau_hip, tau_knee, + max_tau, i2r, imbal, cond, r_hip_x_mag, calf_deg + """ + # --- FK --- + p = mv(ry(q_knee), T_WHEEL) + p = add(p, T_KNEE) + p = mv(ry(q_hip), p) + p = add(p, T_HIP) + p = mv(rx(q_ab), p) + wh = add(p, T_AB) + base_z = 0.10 - wh[2] + wheel_z = wh[2] + + # 膝位置（用于 wheel-below-knee） + pk = mv(ry(q_hip), T_KNEE) + pk = mv(rx(q_ab), pk) + kz = pk[2] + T_AB[2] + + # 关节位置 + jk = mv(ry(q_hip), add(T_KNEE, T_HIP)) + jk = mv(rx(q_ab), jk) + jk = add(T_AB, jk) + jh = mv(rx(q_ab), T_HIP) + jh = add(T_AB, jh) + + # 力矩 + rk = [wh[0]-jk[0], wh[1]-jk[1], wh[2]-jk[2]] + rh = [wh[0]-jh[0], wh[1]-jh[1], wh[2]-jh[2]] + ra = [wh[0]-T_AB[0], wh[1]-T_AB[1], wh[2]-T_AB[2]] + tau_ab = ra[1] * F_PER_LEG + tau_hip = -rh[0] * F_PER_LEG + tau_knee = -rk[0] * F_PER_LEG + + # --- Jacobian（有限差分） --- + eps = 1e-6 + def foot_rel_hip(h, k): + fp = mv(ry(k), T_WHEEL) + fp = add(fp, T_KNEE) + fp = mv(ry(h), fp) + return [fp[0] + T_AB[0] - jh[0], fp[2] + T_AB[2] - jh[2]] + fp0 = foot_rel_hip(q_hip, q_knee) + fph = foot_rel_hip(q_hip + eps, q_knee) + fpk = foot_rel_hip(q_hip, q_knee + eps) + J = np.array([ + [(fph[0]-fp0[0])/eps, (fpk[0]-fp0[0])/eps], + [(fph[1]-fp0[1])/eps, (fpk[1]-fp0[1])/eps], + ]) + s = np.linalg.svd(J, compute_uv=False) + cond = s[0] / s[-1] if s[-1] > 1e-10 else 999.0 + min_sv = s[-1] + + # --- 小腿角度（相对铅垂线） --- + calf_x = (-0.20011) * math.sin(q_knee) + calf_z = (-0.20011) * math.cos(q_knee) + cv_x = calf_x * math.cos(q_hip) + calf_z * math.sin(q_hip) + cv_z = -calf_x * math.sin(q_hip) + calf_z * math.cos(q_hip) + calf_deg = math.degrees(math.atan2(cv_x, -cv_z)) + + abs_taus = [abs(tau_ab), abs(tau_hip), abs(tau_knee)] + return { + 'z': base_z, 'ab': q_ab, 'hip': q_hip, 'knee': q_knee, + 'tau_ab': tau_ab, 'tau_hip': tau_hip, 'tau_knee': tau_knee, + 'max_tau': max(abs_taus), + 'i2r': tau_ab**2 + tau_hip**2 + tau_knee**2, + 'imbal': max(abs_taus) / max(1e-10, min(abs_taus)), + 'cond': cond, 'min_sv': min_sv, + 'r_hip_x_mag': abs(rh[0]), + 'calf_deg': calf_deg, + 'kz': kz, 'wz': wheel_z, + } + + +def rl_cost(r): + """RL友好度综合成本（越小越好）。 + + 约束来源： + c1 — 瓶颈电机发热 τ²/τ_max² 主目标 + c2 — 电机不均衡 > 1.3x (imbal-1.3)² 单电机先超载 + c3 — 髋力臂 < 8cm (0.08 - r_hip) 髋闲置→动态响应差 + c4 — 有效传动比 κ > 3.5 (κ - 3.5) 高刚度→冲击传递大 + c5 — 腿的被动刚度 > 1.3x (stiff-1.3) 刚度比→冲击吸收（新！） + """ + c1 = (r['max_tau'] / 17.0) ** 2 + c2 = max(0.0, (r['imbal'] - 1.3) / 1.0) ** 2 + c3 = max(0.0, (0.08 - r['r_hip_x_mag'])) / 0.08 + c4 = max(0.0, (r['cond'] - 3.5)) / 5.0 + # stiffness ratio normalized to z=0.40 (σ_min≈0.115) + stiff = (0.115 / r['min_sv']) ** 2 + c5 = max(0.0, (stiff - 1.3)) / 3.0 + return 100.0*c1 + 50.0*c2 + 80.0*c3 + 30.0*c4 + 40.0*c5 + + +def sweep_z(z_targets, name, ab_max=0.44, tol=0.004): + """遍历 z 扫描最优姿态。""" + print(f"\n{'='*100}") + print(f" {name}") + print(f"{'='*100}") + print(f"{'z_tgt':>6} {'z':>6} {'ab':>5} {'hip':>6} {'knee':>6} | " + f"{'maxτ':>6} {'imbal':>6} {'r_hip':>6} {'κ':>5} | " + f"{'c1热':>6} {'c2均':>6} {'c3髋':>6} {'c4奇':>6} {'cost':>7}") + print("-"*100) + results = [] + for zt in z_targets: + best_cost = float('inf') + best_r = None + for ab in [round(i*0.02, 2) for i in range(int(ab_max/0.02)+1)]: + for h in [round(i*0.01, 2) for i in range(260)]: + for kn in [round(-2.65+i*0.01, 2) for i in range(256)]: + if h + kn > -0.1: continue + r = get_all(ab, h, kn) + if abs(r['z'] - zt) > tol: continue + if r['wz'] >= r['kz']: continue + if r['tau_ab'] < 0 or r['tau_hip'] < 0: continue + cost = rl_cost(r) + if cost < best_cost: + best_cost = cost; best_r = r + if best_r: + r = best_r + c1s = 100.0*(r['max_tau']/17.0)**2 + c2s = 50.0*max(0.0,(r['imbal']-1.3))**2 + c3s = 80.0*max(0.0,(0.08-r['r_hip_x_mag']))/0.08 + c4s = 30.0*max(0.0,(r['cond']-3.5))/5.0 + print(f"{zt:>6.2f} {r['z']:>6.3f} {r['ab']:>5.2f} {r['hip']:>6.2f} {r['knee']:>6.2f} | " + f"{r['max_tau']:>6.3f} {r['imbal']:>6.1f}x {r['r_hip_x_mag']:>6.3f} {r['cond']:>5.1f} | " + f"{c1s:>6.1f} {c2s:>6.1f} {c3s:>6.1f} {c4s:>6.1f} {best_cost:>7.1f}") + results.append((zt, r)) + else: + print(f"{zt:>6.2f} — no valid") + return results + + +if __name__ == "__main__": + stand = sweep_z([round(0.36+0.01*i, 2) for i in range(10)], + "STANDING — heat + RL constraints (ab free)", 0.44) + crawl = sweep_z([round(0.10+0.01*i, 2) for i in range(9)], + "CRAWL — heat + RL constraints (ab_max=0.25)", 0.25) + + print(f"\n{'='*100}") + print(" RECOMMENDATION") + print(f"{'='*100}") + if stand: + zt, r = min(stand, key=lambda x: rl_cost(x[1])) + print(f" Standing: z={r['z']:.3f} ab={r['ab']:.2f} hip={r['hip']:.2f} knee={r['knee']:.2f}") + print(f" maxτ={r['max_tau']:.3f} imbal={r['imbal']:.1f}x κ={r['cond']:.1f} r_hip={r['r_hip_x_mag']:.3f}m") + if crawl: + zt, r = min(crawl, key=lambda x: rl_cost(x[1])) + print(f" Crawl: z={r['z']:.3f} ab={r['ab']:.2f} hip={r['hip']:.2f} knee={r['knee']:.2f}") + print(f" maxτ={r['max_tau']:.3f} imbal={r['imbal']:.1f}x κ={r['cond']:.1f} r_hip={r['r_hip_x_mag']:.3f}m") diff --git a/05_software/train/rc_mjlab/mujoco_sim/robot.py b/05_software/train/rc_mjlab/mujoco_sim/robot.py new file mode 100644 index 0000000..24af70d --- /dev/null +++ b/05_software/train/rc_mjlab/mujoco_sim/robot.py @@ -0,0 +1,231 @@ +"""MuJoCo interface for the wheeled-legged robot. + +Configures actuators as proper PD controllers at runtime: + - Leg joints: force = kp*(ctrl - qpos) - kd*qvel, ctrl = target angle + - Wheel joints: force = gain*(ctrl - qvel), ctrl = target velocity (rad/s) +""" + +import math +import numpy as np +import mujoco +from dataclasses import dataclass +from config import (SCENE_XML, LEG_NAMES, LEG_JOINTS, WHEEL_JOINT, + DEFAULT_JOINT_ANGLES, WHEEL_RADIUS, WHEEL_TRACK, + L_THIGH, L_CALF) + + +@dataclass +class RobotState: + """Robot state from MuJoCo.""" + pos: np.ndarray # (3,) world position + quat: np.ndarray # (4,) quaternion (w,x,y,z) MuJoCo convention + rot: np.ndarray # (3,3) body→world rotation + rpy: np.ndarray # (3,) roll, pitch, yaw + lin_vel: np.ndarray # (3,) world frame linear velocity + ang_vel: np.ndarray # (3,) body frame angular velocity + joint_pos: np.ndarray # (16,) all joint positions [fl3+wheel, fr3+wheel, rl3+wheel, rr3+wheel] + joint_vel: np.ndarray # (16,) all joint velocities + time: float + + +class Robot: + """MuJoCo simulation interface with proper PD actuator configuration.""" + + # Leg PD gains (tuned for 12.3kg robot) + LEG_KP = 60.0 + LEG_KD = 3.0 + # Wheel velocity gain + WHEEL_KP = 2.0 + + def __init__(self, xml_path=None): + self.model = mujoco.MjModel.from_xml_path(str(xml_path or SCENE_XML)) + self.data = mujoco.MjData(self.model) + + # Cache IDs + self._base_bid = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, "base_link") + self._actuator_ids = {} # name → actuator index + self._joint_qpos_adr = {} # name → qpos address + self._joint_qvel_adr = {} # name → qvel address + + # Build joint/actuator maps + self._ctrl_order = [] + for leg in LEG_NAMES: + for jt in (*LEG_JOINTS, WHEEL_JOINT): + name = f"{leg}_{jt}" + aid = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_ACTUATOR, name) + jid = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, name) + self._actuator_ids[name] = aid + self._joint_qpos_adr[name] = self.model.jnt_qposadr[jid] + self._joint_qvel_adr[name] = self.model.jnt_dofadr[jid] + self._ctrl_order.append(name) + + # Configure actuators as proper PD controllers + self._configure_actuators() + + def _configure_actuators(self): + """Set actuators to proper PD mode. + + Leg joints: force = kp*(ctrl - qpos) - kd*qvel + Wheels: force = gain*(ctrl - qvel) (velocity tracking) + """ + for i in range(self.model.nu): + name = mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_ACTUATOR, i) + self.model.actuator_biastype[i] = 1 # affine bias + self.model.actuator_gaintype[i] = 0 # fixed gain + self.model.actuator_forcelimited[i] = 0 # no force clamp (17Nm is in actuatorfrcrange) + + if 'wheel' not in name: + self.model.actuator_gainprm[i, 0] = self.LEG_KP + self.model.actuator_biasprm[i, 0] = 0.0 + self.model.actuator_biasprm[i, 1] = -self.LEG_KP + self.model.actuator_biasprm[i, 2] = -self.LEG_KD + self.model.actuator_ctrlrange[i] = [-3.14, 3.14] + else: + self.model.actuator_gainprm[i, 0] = self.WHEEL_KP + self.model.actuator_biasprm[i, 0] = 0.0 + self.model.actuator_biasprm[i, 1] = 0.0 + self.model.actuator_biasprm[i, 2] = -self.WHEEL_KP + self.model.actuator_ctrlrange[i] = [-20.0, 20.0] + + @property + def dt(self): + return self.model.opt.timestep + + def reset(self): + """Reset to standing pose at correct height for default joint angles.""" + mujoco.mj_resetData(self.model, self.data) + + # Set default leg angles + for leg in LEG_NAMES: + for jt, key in zip(LEG_JOINTS, ("hip_abduction", "hip_pitch", "knee")): + name = f"{leg}_{jt}" + adr = self._joint_qpos_adr[name] + self.data.qpos[adr] = DEFAULT_JOINT_ANGLES[key] + + # Compute correct base height from default angles using exact FK + # z_base = wheel_radius + L1*cos(hip) + L2*cos(hip+knee) - hip_z_offset + hip = DEFAULT_JOINT_ANGLES["hip_pitch"] + knee = DEFAULT_JOINT_ANGLES["knee"] + base_z = (WHEEL_RADIUS + + L_THIGH * math.cos(hip) + + L_CALF * math.cos(hip + knee) + - 0.054) + self.data.qpos[2] = base_z + self.data.qpos[3] = 1.0 # quat w + + mujoco.mj_forward(self.model, self.data) + + # Set ctrl to match initial pose (so PD doesn't jerk) + for leg in LEG_NAMES: + for jt, key in zip(LEG_JOINTS, ("hip_abduction", "hip_pitch", "knee")): + name = f"{leg}_{jt}" + self.data.ctrl[self._actuator_ids[name]] = DEFAULT_JOINT_ANGLES[key] + # Wheels: zero velocity + self.data.ctrl[self._actuator_ids[f"{leg}_{WHEEL_JOINT}"]] = 0.0 + + def get_state(self) -> RobotState: + """Extract robot state.""" + pos = self.data.xpos[self._base_bid].copy() + quat = self.data.xquat[self._base_bid].copy() # (w,x,y,z) + rot = self.data.xmat[self._base_bid].reshape(3, 3).copy() + + rpy = np.array([ + np.arctan2(rot[2, 1], rot[2, 2]), + np.arctan2(-rot[2, 0], np.sqrt(rot[2, 1]**2 + rot[2, 2]**2)), + np.arctan2(rot[1, 0], rot[0, 0]), + ]) + + # Base velocity (world frame) + lin_vel = self.data.qvel[0:3].copy() + ang_vel = self.data.qvel[3:6].copy() + + # Joint states (16 joints: 4 legs × 4 joints each) + joint_pos = np.zeros(16) + joint_vel = np.zeros(16) + for i, name in enumerate(self._ctrl_order): + joint_pos[i] = self.data.qpos[self._joint_qpos_adr[name]] + joint_vel[i] = self.data.qvel[self._joint_qvel_adr[name]] + + return RobotState( + pos=pos, quat=quat, rot=rot, rpy=rpy, + lin_vel=lin_vel, ang_vel=ang_vel, + joint_pos=joint_pos, joint_vel=joint_vel, + time=self.data.time, + ) + + def set_ctrl(self, leg_targets: np.ndarray, wheel_targets: np.ndarray): + """Set actuator commands (position PD mode). + + Args: + leg_targets: (12,) target joint angles for legs [fl3, fr3, rl3, rr3] + wheel_targets: (4,) target wheel velocities [fl, fr, rl, rr] in rad/s + """ + for i, leg in enumerate(LEG_NAMES): + for j, jt in enumerate(LEG_JOINTS): + name = f"{leg}_{jt}" + self.data.ctrl[self._actuator_ids[name]] = leg_targets[i * 3 + j] + name = f"{leg}_{WHEEL_JOINT}" + self.data.ctrl[self._actuator_ids[name]] = wheel_targets[i] + + def set_ctrl_mit(self, q_des: np.ndarray, dq_des: np.ndarray, + kp: np.ndarray, kd: np.ndarray, tau_ff: np.ndarray, + wheel_targets: np.ndarray): + """MIT motor protocol: tau = kp*(q_des-q) + kd*(dq_des-dq) + tau_ff. + + Computes torque in software, sends to actuators in torque mode. + Call enable_torque_mode() first. + + Args: + q_des: (12,) desired joint angles + dq_des: (12,) desired joint velocities + kp: (12,) position gains (0 for pure torque) + kd: (12,) velocity gains + tau_ff: (12,) feedforward torques + wheel_targets: (4,) wheel velocity targets + """ + for i, leg in enumerate(LEG_NAMES): + for j, jt in enumerate(LEG_JOINTS): + name = f"{leg}_{jt}" + aid = self._actuator_ids[name] + idx = i * 3 + j + q = self.data.qpos[self._joint_qpos_adr[name]] + dq = self.data.qvel[self._joint_qvel_adr[name]] + tau = (kp[idx] * (q_des[idx] - q) + + kd[idx] * (dq_des[idx] - dq) + + tau_ff[idx]) + self.data.ctrl[aid] = np.clip(tau, -17.0, 17.0) + name = f"{leg}_{WHEEL_JOINT}" + self.data.ctrl[self._actuator_ids[name]] = wheel_targets[i] + + def enable_torque_mode(self): + """Switch leg actuators to direct torque mode (for MPC/MIT).""" + for i in range(self.model.nu): + name = mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_ACTUATOR, i) + if 'wheel' not in name: + self.model.actuator_gainprm[i, 0] = 1.0 + self.model.actuator_biasprm[i, :3] = [0, 0, 0] + self.model.actuator_biastype[i] = 0 + self.model.actuator_ctrlrange[i] = [-17.0, 17.0] + + def enable_position_mode(self): + """Switch leg actuators back to position PD mode.""" + self._configure_actuators() + + def step(self): + """Advance one simulation timestep.""" + mujoco.mj_step(self.model, self.data) + + def get_qpos_qvel_for_pinocchio(self): + """Get full qpos/qvel for Pinocchio (reorder quaternion).""" + qpos = self.data.qpos.copy() + qvel = self.data.qvel.copy() + # MuJoCo quat: (w,x,y,z) → Pinocchio: (x,y,z,w) + w, x, y, z = qpos[3], qpos[4], qpos[5], qpos[6] + q_pin = np.concatenate([qpos[0:3], [x, y, z, w], qpos[7:]]) + # MuJoCo vel is already [lin_world(3), ang_body(3), joints(16)] + # Pinocchio wants [lin_body(3), ang_body(3), joints(16)] + from scipy.spatial.transform import Rotation + R = Rotation.from_quat([x, y, z, w]).as_matrix() + v_body = R.T @ qvel[0:3] + dq_pin = np.concatenate([v_body, qvel[3:]]) + return q_pin, dq_pin diff --git a/05_software/train/rc_mjlab/mujoco_sim/run.py b/05_software/train/rc_mjlab/mujoco_sim/run.py new file mode 100644 index 0000000..5fdcdb5 --- /dev/null +++ b/05_software/train/rc_mjlab/mujoco_sim/run.py @@ -0,0 +1,111 @@ +"""Main entry point: wheeled-legged robot simulation. + +Controls: + Mode: wheel (default) - differential drive + posture hold + trot - quadruped gait with wheel assist + mpc - convex MPC locomotion (torque control) + + Keyboard (in MuJoCo viewer): + W/S: vel_x ±0.1 + A/D: yaw_rate ±0.2 + Q/E: height ±0.02 + 1: wheel mode + 2: trot mode + 3: MPC mode + 4: prone toggle + Z: reset commands +""" + +import time +import numpy as np +import mujoco.viewer as mjv + +from robot import Robot +from controller import Controller +from gui import GUI +from config import CTRL_DECIMATION + + +def main(): + robot = Robot() + robot.reset() + ctrl = Controller(robot) + gui = GUI(ctrl) + + step = 0 + sim_steps_per_ctrl = CTRL_DECIMATION + + def key_callback(keycode): + """Called from MuJoCo render thread - only modify ctrl directly, not tkinter.""" + try: + c = chr(keycode).lower() + except (ValueError, OverflowError): + return + if c == 'w': + ctrl.vel_x = min(ctrl.vel_x + 0.1, 1.5) + elif c == 's': + ctrl.vel_x = max(ctrl.vel_x - 0.1, -1.5) + elif c == 'a': + ctrl.yaw_rate = min(ctrl.yaw_rate + 0.2, 2.0) + elif c == 'd': + ctrl.yaw_rate = max(ctrl.yaw_rate - 0.2, -2.0) + elif c == 'q': + ctrl.height = min(ctrl.height + 0.02, 0.45) + elif c == 'e': + ctrl.height = max(ctrl.height - 0.02, 0.17) + elif c == '1': + ctrl.mode = "wheel"; ctrl.prone = False + elif c == '2': + ctrl.mode = "trot"; ctrl.prone = False + elif c == '3': + ctrl.mode = "mpc"; ctrl.prone = False + elif c == '4': + ctrl.prone = not ctrl.prone + elif c == 'z': + ctrl.vel_x = 0.0; ctrl.vel_y = 0.0; ctrl.yaw_rate = 0.0 + + with mjv.launch_passive(robot.model, robot.data, key_callback=key_callback) as viewer: + viewer.cam.distance = 2.5 + viewer.cam.elevation = -20 + viewer.cam.azimuth = 135 + + last_time = robot.data.time + + while viewer.is_running() and not gui.closed: + t_start = time.perf_counter() + + # Detect viewer reset (Backspace) - time jumps back to 0 + if robot.data.time < last_time: + robot.reset() + last_time = robot.data.time + + # Get state and compute control + state = robot.get_state() + leg_targets, wheel_targets = ctrl.compute(state, robot.dt * sim_steps_per_ctrl) + + # Apply control and step simulation + # MPC mode sets ctrl directly via set_ctrl_mit, skip set_ctrl + if ctrl.mode != "mpc": + robot.set_ctrl(leg_targets, wheel_targets) + for _ in range(sim_steps_per_ctrl): + robot.step() + + viewer.sync() + step += 1 + + # Update GUI every 25 steps (~10 Hz) + if step % 25 == 0: + state = robot.get_state() + gui.update_status(state, step) + if not gui.tick(): + break + + # Real-time sync + elapsed = time.perf_counter() - t_start + target_dt = robot.dt * sim_steps_per_ctrl + if elapsed < target_dt: + time.sleep(target_dt - elapsed) + + +if __name__ == "__main__": + main() diff --git a/05_software/train/rc_mjlab/mujoco_sim/static_posture_optimizer.py b/05_software/train/rc_mjlab/mujoco_sim/static_posture_optimizer.py new file mode 100644 index 0000000..73f5c5e --- /dev/null +++ b/05_software/train/rc_mjlab/mujoco_sim/static_posture_optimizer.py @@ -0,0 +1,528 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +"""MuJoCo static posture optimizer for wheeled-leg standing defaults. + +The old posture tools are mostly analytical. This script keeps the fast MJCF +kinematics as a candidate generator, then evaluates the best candidates in +MuJoCo with gravity and floor contact enabled. + +The score is meant for RL default pose / real deployment: + - low peak and RMS standing torque, so one hot motor is not hidden by average + - wheel contact point close to the hip in X for wheel speed tracking + - COM projection margin inside the four-wheel support rectangle + - non-singular leg Jacobian for posture control authority + - underbody and knee clearance for obstacle tolerance + +Usage: + uv run python mujoco_sim/static_posture_optimizer.py + uv run python mujoco_sim/static_posture_optimizer.py --quick + uv run python mujoco_sim/static_posture_optimizer.py --ab-max 0.08 +""" + +from __future__ import annotations + +import argparse +import math +import sys +from dataclasses import dataclass +from pathlib import Path + +import mujoco +import numpy as np + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from mujoco_sim.rl_friendly_opt import get_all # noqa: E402 + +SCENE_XML = REPO_ROOT / "mjcf" / "scene.xml" + +LEG_NAMES = ("fl", "fr", "rl", "rr") +LEG_JOINTS = ("hip_abduction_joint", "hip_pitch_joint", "knee_joint") +WHEEL_JOINT = "wheel_joint" +LEG_ACTUATORS = [f"{leg}_{jt}" for leg in LEG_NAMES for jt in LEG_JOINTS] + +ROBOT_MASS = 12.3 +G = 9.81 +MAX_TORQUE = 17.0 +WHEEL_RADIUS = 0.10 + +HIP_MIN, HIP_MAX = -2.58, 2.58 +KNEE_MIN, KNEE_MAX = -2.65, 2.65 +HIP_SCAN = (0.20, 1.15) +KNEE_SCAN = (-1.90, -0.65) + +SOFT_WHEEL_X = 0.045 +HARD_WHEEL_X = 0.085 +MIN_COM_MARGIN = 0.045 +MIN_KNEE_CLEARANCE = 0.105 +MIN_UNDERBODY_CLEARANCE = 0.33 +MAX_COND = 3.7 + + +@dataclass(frozen=True) +class Candidate: + ab: float + hip: float + knee: float + z_fk: float + analytic_cost: float + + +@dataclass +class StaticResult: + z_target: float + z: float + ab: float + hip: float + knee: float + cost: float + peak_tau: float + rms_tau: float + mean_i2r: float + imbal: float + wheel_x: float + cond: float + min_sv: float + com_margin_x: float + com_margin_y: float + support_margin: float + normal_cv: float + body_clearance: float + knee_clearance: float + roll: float + pitch: float + height_err: float + + +def _configure_actuators(model: mujoco.MjModel) -> None: + """Configure leg joints as position PD and wheels as zero-velocity motors.""" + for i in range(model.nu): + name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_ACTUATOR, i) or "" + model.actuator_biastype[i] = 1 + model.actuator_gaintype[i] = 0 + model.actuator_forcelimited[i] = 0 + if "wheel" in name: + model.actuator_gainprm[i, 0] = 2.0 + model.actuator_biasprm[i, 0] = 0.0 + model.actuator_biasprm[i, 1] = 0.0 + model.actuator_biasprm[i, 2] = -2.0 + model.actuator_ctrlrange[i] = [-20.0, 20.0] + else: + model.actuator_gainprm[i, 0] = 60.0 + model.actuator_biasprm[i, 0] = 0.0 + model.actuator_biasprm[i, 1] = -60.0 + model.actuator_biasprm[i, 2] = -3.0 + model.actuator_ctrlrange[i] = [-3.14, 3.14] + + +def _ids(model: mujoco.MjModel): + act = {} + qadr = {} + for leg in LEG_NAMES: + for jt in (*LEG_JOINTS, WHEEL_JOINT): + name = f"{leg}_{jt}" + act[name] = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, name) + jid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, name) + qadr[name] = model.jnt_qposadr[jid] + bodies = { + "base": mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "base_link"), + **{ + f"{leg}_wheel": mujoco.mj_name2id( + model, mujoco.mjtObj.mjOBJ_BODY, f"{leg}_wheel_Link" + ) + for leg in LEG_NAMES + }, + **{ + f"{leg}_knee": mujoco.mj_name2id( + model, mujoco.mjtObj.mjOBJ_BODY, f"{leg}_knee_Link" + ) + for leg in LEG_NAMES + }, + } + return act, qadr, bodies + + +def _mirrored_ab(leg: str, ab: float) -> float: + return ab if leg[1] == "l" else -ab + + +def _set_pose( + model: mujoco.MjModel, + data: mujoco.MjData, + act: dict[str, int], + qadr: dict[str, int], + cand: Candidate, +) -> None: + mujoco.mj_resetData(model, data) + for leg in LEG_NAMES: + vals = (_mirrored_ab(leg, cand.ab), cand.hip, cand.knee) + for jt, val in zip(LEG_JOINTS, vals): + name = f"{leg}_{jt}" + data.qpos[qadr[name]] = val + data.ctrl[act[name]] = val + data.ctrl[act[f"{leg}_{WHEEL_JOINT}"]] = 0.0 + data.qpos[0:3] = [0.0, 0.0, cand.z_fk] + data.qpos[3] = 1.0 + data.qpos[4:7] = 0.0 + mujoco.mj_forward(model, data) + + +def _rpy(data: mujoco.MjData, bodies: dict[str, int]) -> tuple[float, float]: + rot = data.xmat[bodies["base"]].reshape(3, 3) + roll = math.atan2(rot[2, 1], rot[2, 2]) + pitch = math.atan2(-rot[2, 0], math.sqrt(rot[2, 1] ** 2 + rot[2, 2] ** 2)) + return roll, pitch + + +def _robot_com(model: mujoco.MjModel, data: mujoco.MjData) -> np.ndarray: + masses = model.body_mass[1:] + return (data.xipos[1:] * masses[:, None]).sum(axis=0) / masses.sum() + + +def _support_metrics( + model: mujoco.MjModel, + data: mujoco.MjData, + bodies: dict[str, int], +) -> tuple[float, float, float, float]: + wheel_xy = np.array([data.xpos[bodies[f"{leg}_wheel"]][:2] for leg in LEG_NAMES]) + com_xy = _robot_com(model, data)[:2] + min_xy = wheel_xy.min(axis=0) + max_xy = wheel_xy.max(axis=0) + margin_low = com_xy - min_xy + margin_high = max_xy - com_xy + margin_x = float(min(margin_low[0], margin_high[0])) + margin_y = float(min(margin_low[1], margin_high[1])) + support_margin = float(min(margin_x, margin_y)) + + normal = [] + for leg in LEG_NAMES: + bid = bodies[f"{leg}_wheel"] + fz = 0.0 + for i in range(data.ncon): + con = data.contact[i] + b1 = model.geom_bodyid[con.geom1] + b2 = model.geom_bodyid[con.geom2] + if b1 == bid or b2 == bid: + wrench = np.zeros(6) + mujoco.mj_contactForce(model, data, i, wrench) + fz += abs(float(wrench[0])) + normal.append(fz) + normal = np.asarray(normal, dtype=float) + if normal.sum() < 1e-6: + normal_cv = 9.99 + else: + normal_cv = float(normal.std() / max(1e-6, normal.mean())) + return margin_x, margin_y, support_margin, normal_cv + + +def _clearance_metrics(data: mujoco.MjData, bodies: dict[str, int]) -> tuple[float, float]: + base_z = float(data.xpos[bodies["base"]][2]) + # The collision box in wheelleg.xml is centered at z=0.054 with half-height 0.073. + body_clearance = base_z + 0.054 - 0.073 + knee_z = min(float(data.xpos[bodies[f"{leg}_knee"]][2]) for leg in LEG_NAMES) + wheel_z = min(float(data.xpos[bodies[f"{leg}_wheel"]][2]) for leg in LEG_NAMES) + return body_clearance, knee_z - wheel_z + + +def _static_cost(r: StaticResult) -> float: + wheel_over = max(0.0, abs(r.wheel_x) - SOFT_WHEEL_X) + support_short = max(0.0, MIN_COM_MARGIN - r.support_margin) + body_short = max(0.0, MIN_UNDERBODY_CLEARANCE - r.body_clearance) + knee_short = max(0.0, MIN_KNEE_CLEARANCE - r.knee_clearance) + cond_over = max(0.0, r.cond - MAX_COND) + tilt = math.hypot(r.roll, r.pitch) + + return ( + 6500.0 * r.height_err**2 + + 2.6 * (r.peak_tau / MAX_TORQUE) ** 2 + + 1.1 * (r.rms_tau / MAX_TORQUE) ** 2 + + 1.8 * max(0.0, r.imbal - 1.6) ** 2 + + 0.65 * (abs(r.wheel_x) / SOFT_WHEEL_X) ** 2 + + 9.0 * (wheel_over / max(1e-6, HARD_WHEEL_X - SOFT_WHEEL_X)) ** 2 + + 10.0 * (support_short / MIN_COM_MARGIN) ** 2 + + 2.0 * (r.normal_cv / 0.35) ** 2 + + 30.0 * (cond_over / 1.0) ** 2 + + 4.0 * (body_short / 0.06) ** 2 + + 2.0 * (knee_short / 0.04) ** 2 + + 1.0 * (tilt / 0.05) ** 2 + ) + + +def _analytic_cost(r: dict, z_target: float) -> float: + wheel_over = max(0.0, abs(r["r_hip_x_mag"]) - SOFT_WHEEL_X) + cond_over = max(0.0, r["cond"] - MAX_COND) + return ( + 1400.0 * (r["z"] - z_target) ** 2 + + 2.2 * (r["max_tau"] / MAX_TORQUE) ** 2 + + 0.7 * (math.sqrt(r["i2r"] / 3.0) / MAX_TORQUE) ** 2 + + 0.55 * (abs(r["r_hip_x_mag"]) / SOFT_WHEEL_X) ** 2 + + 8.0 * (wheel_over / max(1e-6, HARD_WHEEL_X - SOFT_WHEEL_X)) ** 2 + + 2.0 * (cond_over / 2.0) ** 2 + ) + + +def generate_candidates( + z_target: float, + step: float, + ab_max: float, + keep: int, + z_tol: float, +) -> list[Candidate]: + candidates: list[Candidate] = [] + ab_values = np.arange(0.0, ab_max + 0.5 * step, step) + hip_values = np.arange(HIP_SCAN[0], HIP_SCAN[1] + 0.5 * step, step) + knee_values = np.arange(KNEE_SCAN[0], KNEE_SCAN[1] + 0.5 * step, step) + for ab in ab_values: + for hip in hip_values: + if hip < HIP_MIN or hip > HIP_MAX: + continue + for knee in knee_values: + if knee < KNEE_MIN or knee > KNEE_MAX: + continue + r = get_all(float(ab), float(hip), float(knee)) + if abs(r["z"] - z_target) > z_tol: + continue + if r["wz"] >= r["kz"]: + continue + if abs(r["r_hip_x_mag"]) > HARD_WHEEL_X: + continue + if r["max_tau"] > MAX_TORQUE * 1.15: + continue + candidates.append( + Candidate( + ab=float(ab), + hip=float(hip), + knee=float(knee), + z_fk=float(r["z"]), + analytic_cost=_analytic_cost(r, z_target), + ) + ) + candidates.sort(key=lambda c: c.analytic_cost) + return candidates[:keep] + + +def evaluate_candidate( + model: mujoco.MjModel, + data: mujoco.MjData, + act: dict[str, int], + qadr: dict[str, int], + bodies: dict[str, int], + cand: Candidate, + z_target: float, + settle_steps: int, + avg_steps: int, +) -> StaticResult | None: + _set_pose(model, data, act, qadr, cand) + for _ in range(settle_steps): + mujoco.mj_step(model, data) + + roll, pitch = _rpy(data, bodies) + if abs(roll) > 0.35 or abs(pitch) > 0.35: + return None + + tau_buf = [] + for _ in range(avg_steps): + mujoco.mj_step(model, data) + tau_buf.append([data.actuator_force[act[name]] for name in LEG_ACTUATORS]) + tau = np.asarray(tau_buf, dtype=float).mean(axis=0) + tau_abs = np.abs(tau) + peak = float(tau_abs.max()) + rms = float(np.sqrt(np.mean(tau * tau))) + mean_i2r = float(np.mean(tau * tau)) + + sagittal_abs = [] + for name, value in zip(LEG_ACTUATORS, tau): + if "hip_pitch" in name or "knee" in name: + sagittal_abs.append(abs(float(value))) + sagittal_abs = np.asarray(sagittal_abs, dtype=float) + imbal = float(sagittal_abs.max() / max(1e-6, sagittal_abs.mean())) + + fk = get_all(cand.ab, cand.hip, cand.knee) + margin_x, margin_y, support_margin, normal_cv = _support_metrics(model, data, bodies) + body_clearance, knee_clearance = _clearance_metrics(data, bodies) + z = float(data.xpos[bodies["base"]][2]) + + result = StaticResult( + z_target=z_target, + z=z, + ab=cand.ab, + hip=cand.hip, + knee=cand.knee, + cost=0.0, + peak_tau=peak, + rms_tau=rms, + mean_i2r=mean_i2r, + imbal=imbal, + wheel_x=float(fk["r_hip_x_mag"]), + cond=float(fk["cond"]), + min_sv=float(fk["min_sv"]), + com_margin_x=margin_x, + com_margin_y=margin_y, + support_margin=support_margin, + normal_cv=normal_cv, + body_clearance=body_clearance, + knee_clearance=knee_clearance, + roll=roll, + pitch=pitch, + height_err=z - z_target, + ) + result.cost = _static_cost(result) + return result + + +def optimize( + z_targets: list[float], + step: float, + ab_max: float, + keep: int, + z_tol: float, + settle_steps: int, + avg_steps: int, +) -> list[StaticResult]: + model = mujoco.MjModel.from_xml_path(str(SCENE_XML)) + data = mujoco.MjData(model) + _configure_actuators(model) + act, qadr, bodies = _ids(model) + + results: list[StaticResult] = [] + for zt in z_targets: + candidates = generate_candidates(zt, step=step, ab_max=ab_max, keep=keep, z_tol=z_tol) + best: StaticResult | None = None + for cand in candidates: + result = evaluate_candidate( + model, + data, + act, + qadr, + bodies, + cand, + zt, + settle_steps=settle_steps, + avg_steps=avg_steps, + ) + if result is None: + continue + if best is None or result.cost < best.cost: + best = result + if best is None: + print(f"{zt:.2f}: no valid MuJoCo-static candidate from {len(candidates)} seeds") + continue + results.append(best) + print_result(best) + return results + + +def print_header() -> None: + print("\nMuJoCo static posture optimization") + print("z_tgt z ab hip knee | peak rms imbal wheelX cond | comX comY clrB clrK | cost") + print("-" * 111) + + +def print_result(r: StaticResult) -> None: + print( + f"{r.z_target:5.2f} {r.z:6.3f} {r.ab:5.2f} {r.hip:6.3f} {r.knee:6.3f} | " + f"{r.peak_tau:5.2f} {r.rms_tau:5.2f} {r.imbal:6.2f} " + f"{r.wheel_x:6.3f} {r.cond:5.2f} | " + f"{r.com_margin_x:5.3f} {r.com_margin_y:5.3f} " + f"{r.body_clearance:5.3f} {r.knee_clearance:5.3f} | " + f"{r.cost:6.2f}" + ) + + +def print_tables(results: list[StaticResult]) -> None: + if not results: + return + print("\nCopy-paste tables:") + print("_H_TARGET = " + repr([round(r.z_target, 3) for r in results])) + print("_Z_STATIC = " + repr([round(r.z, 3) for r in results])) + print("_HIP = " + repr([round(r.hip, 3) for r in results])) + print("_KNEE = " + repr([round(r.knee, 3) for r in results])) + print("_ABD_LEFT = " + repr([round(r.ab, 3) for r in results])) + print("_ABD_RIGHT = " + repr([round(-r.ab, 3) for r in results])) + + best = min(results, key=lambda r: r.cost) + print("\nRecommended default:") + print( + f"z={best.z_target:.2f}, ab={best.ab:.3f}, hip={best.hip:.3f}, " + f"knee={best.knee:.3f}, peak={best.peak_tau:.2f}Nm, " + f"rms={best.rms_tau:.2f}Nm, support_margin={best.support_margin:.3f}m" + ) + pose = [] + for leg in LEG_NAMES: + pose.extend([_mirrored_ab(leg, best.ab), best.hip, best.knee]) + pose.extend([0.0, 0.0, 0.0, 0.0]) + print("default_dof_pos = " + repr([round(v, 3) for v in pose])) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--z-min", type=float, default=0.36) + parser.add_argument("--z-max", type=float, default=0.45) + parser.add_argument("--z-step", type=float, default=0.01) + parser.add_argument("--grid-step", type=float, default=0.01) + parser.add_argument("--ab-max", type=float, default=0.0) + parser.add_argument("--keep", type=int, default=45) + parser.add_argument("--z-tol", type=float, default=0.035) + parser.add_argument("--settle-steps", type=int, default=350) + parser.add_argument("--avg-steps", type=int, default=80) + parser.add_argument("--quick", action="store_true", help="Coarser and faster scan") + args = parser.parse_args() + + if args.quick: + args.grid_step = max(args.grid_step, 0.02) + args.keep = min(args.keep, 20) + args.z_tol = max(args.z_tol, 0.040) + args.settle_steps = min(args.settle_steps, 220) + args.avg_steps = min(args.avg_steps, 40) + + n = int(round((args.z_max - args.z_min) / args.z_step)) + 1 + z_targets = [round(args.z_min + i * args.z_step, 3) for i in range(n)] + + print_header() + results = optimize( + z_targets, + step=args.grid_step, + ab_max=args.ab_max, + keep=args.keep, + z_tol=args.z_tol, + settle_steps=args.settle_steps, + avg_steps=args.avg_steps, + ) + print_tables(results) + + +if __name__ == "__main__": + main() diff --git a/README.md b/README.md index 4e6e208..e6b0c26 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ RC_WheelLeg 是山东华宇工学院 16DOF 串联轮足机器人项目。 -当前 `16dof` 分支用于整理 16DOF 机械、强化学习训练、Sim2Sim、Sim2Real、ROS 2 部署和比赛版本。机械资料、第一代软件闭环和比赛最终训练架构已经完成整理。 +当前 `16dof` 分支用于整理 16DOF 机械、强化学习训练、Sim2Sim、Sim2Real、ROS 2 部署和比赛版本。机械资料、比赛训练架构和后期 MuJoCo 工具集已经完成整理。 ## 平台概览 @@ -36,6 +36,7 @@ RC_WheelLeg/ - [x] 整理第一份新版 MJCF 与 mjlab 训练框架 - [x] 整理第二版 Sim2Real 随机化训练配置 - [x] 整理比赛最终训练代码架构 +- [x] 整理后期 MuJoCo 姿态、IK、动力学和 MPC 工具 - [ ] 核对比赛机械与仿真模型参数 - [ ] 整理 URDF/MJCF 机器人描述 - [ ] 整理后续统一训练、ROS 2 和比赛版本