|
|
@@ -0,0 +1,395 @@
|
|
|
+from __future__ import annotations
|
|
|
+
|
|
|
+import math
|
|
|
+from typing import Any
|
|
|
+
|
|
|
+import gymnasium as gym
|
|
|
+import mujoco
|
|
|
+import numpy as np
|
|
|
+
|
|
|
+from ..config import resolve_project_path
|
|
|
+from ..math_utils import quaternion_xyzw_to_euler
|
|
|
+from ..mujoco_model import build_mujoco_model_spec
|
|
|
+from ..rewards import BipedRewardCalculator, RewardContext
|
|
|
+from ..state_types import RobotStateSnapshot
|
|
|
+from ..urdf_utils import JointLimit, parse_joint_limits
|
|
|
+
|
|
|
+
|
|
|
+class MujocoBipedEnv(gym.Env):
|
|
|
+ """把 guguji 机器人封装成 MuJoCo + Gymnasium 训练环境。"""
|
|
|
+
|
|
|
+ metadata = {'render_modes': ['human', 'rgb_array'], 'render_fps': 20}
|
|
|
+
|
|
|
+ def __init__(self, config: dict[str, Any]) -> None:
|
|
|
+ super().__init__()
|
|
|
+ self.config = config
|
|
|
+ self.robot_config = config['robot']
|
|
|
+ self.sim_config = config['sim']
|
|
|
+ self.task_config = config['task']
|
|
|
+ self.training_config = config['training']
|
|
|
+ self.mujoco_config = config.get('mujoco', {})
|
|
|
+ self.configured_target_base_height = self.task_config['target_base_height']
|
|
|
+ self.reference_gait_config = self.robot_config.get('reference_gait', {})
|
|
|
+
|
|
|
+ urdf_path = resolve_project_path(config, self.robot_config['urdf_path'])
|
|
|
+ self.joint_limits: list[JointLimit] = parse_joint_limits(urdf_path, self.robot_config['joint_names'])
|
|
|
+ self.joint_names = [joint_limit.name for joint_limit in self.joint_limits]
|
|
|
+ self.joint_lower = np.array([joint.lower for joint in self.joint_limits], dtype=np.float32)
|
|
|
+ self.joint_upper = np.array([joint.upper for joint in self.joint_limits], dtype=np.float32)
|
|
|
+ self.joint_mid = np.array([joint.midpoint for joint in self.joint_limits], dtype=np.float32)
|
|
|
+ self.joint_half_range = np.array([joint.half_range for joint in self.joint_limits], dtype=np.float32)
|
|
|
+
|
|
|
+ nominal_joint_config = self.robot_config.get('nominal_joint_positions', {})
|
|
|
+ self.nominal_joint_targets = np.array(
|
|
|
+ [float(nominal_joint_config.get(joint_name, 0.0)) for joint_name in self.joint_names],
|
|
|
+ dtype=np.float32,
|
|
|
+ )
|
|
|
+ self.action_scale = float(self.robot_config.get('action_scale', 1.0))
|
|
|
+ self.action_smoothing = float(np.clip(self.robot_config.get('action_smoothing', 0.0), 0.0, 0.99))
|
|
|
+ self.reference_gait_enabled = bool(self.reference_gait_config.get('enabled', False))
|
|
|
+ self.reference_gait_period = max(
|
|
|
+ float(self.reference_gait_config.get('period', self.sim_config.get('control_dt', 0.05))),
|
|
|
+ 1e-4,
|
|
|
+ )
|
|
|
+ self.termination_grace_steps = int(self.task_config.get('termination_grace_steps', 0))
|
|
|
+
|
|
|
+ self.control_dt = float(self.sim_config.get('control_dt', 0.05))
|
|
|
+ self.frame_skip = int(self.mujoco_config.get('frame_skip', max(int(round(self.control_dt / 0.005)), 1)))
|
|
|
+ self.render_mode = str(self.mujoco_config.get('render_mode', 'none')).lower()
|
|
|
+ if self.render_mode == 'none':
|
|
|
+ self.render_mode = None
|
|
|
+
|
|
|
+ spec = build_mujoco_model_spec(config)
|
|
|
+ self.model = mujoco.MjModel.from_xml_string(spec.xml, spec.assets)
|
|
|
+ self.data = mujoco.MjData(self.model)
|
|
|
+ self.root_body_name = spec.root_link_name
|
|
|
+ self.root_body_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, self.root_body_name)
|
|
|
+
|
|
|
+ self.joint_qpos_indices = np.array(
|
|
|
+ [
|
|
|
+ int(self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, joint_name)])
|
|
|
+ for joint_name in self.joint_names
|
|
|
+ ],
|
|
|
+ dtype=np.int32,
|
|
|
+ )
|
|
|
+ self.joint_qvel_indices = np.array(
|
|
|
+ [
|
|
|
+ int(self.model.jnt_dofadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, joint_name)])
|
|
|
+ for joint_name in self.joint_names
|
|
|
+ ],
|
|
|
+ dtype=np.int32,
|
|
|
+ )
|
|
|
+
|
|
|
+ self.reward_calculator = BipedRewardCalculator(config['rewards'])
|
|
|
+ self.action_space = gym.spaces.Box(
|
|
|
+ low=-1.0,
|
|
|
+ high=1.0,
|
|
|
+ shape=(len(self.joint_names),),
|
|
|
+ dtype=np.float32,
|
|
|
+ )
|
|
|
+ self.observation_space = gym.spaces.Box(
|
|
|
+ low=-np.inf,
|
|
|
+ high=np.inf,
|
|
|
+ shape=(len(self.joint_names) * 3 + 5,),
|
|
|
+ dtype=np.float32,
|
|
|
+ )
|
|
|
+
|
|
|
+ self.base_height = float(self.mujoco_config.get('root_body_height', 0.36))
|
|
|
+ self.base_xy_noise_scale = float(self.mujoco_config.get('reset_base_xy_noise_scale', 0.005))
|
|
|
+ self.base_height_noise_scale = float(self.mujoco_config.get('reset_base_height_noise_scale', 0.002))
|
|
|
+ self.joint_reset_noise_scale = float(self.mujoco_config.get('reset_joint_noise_scale', 0.01))
|
|
|
+ self.reset_hold_steps = int(self.mujoco_config.get('reset_hold_steps', 8))
|
|
|
+ self.render_width = int(self.mujoco_config.get('render_width', 1280))
|
|
|
+ self.render_height = int(self.mujoco_config.get('render_height', 720))
|
|
|
+
|
|
|
+ self.previous_action = np.zeros(len(self.joint_names), dtype=np.float32)
|
|
|
+ self.previous_snapshot: RobotStateSnapshot | None = None
|
|
|
+ self.current_snapshot: RobotStateSnapshot | None = None
|
|
|
+ self.current_joint_target = self.nominal_joint_targets.copy()
|
|
|
+ self.current_action_residual = np.zeros(len(self.joint_names), dtype=np.float32)
|
|
|
+ self.target_base_height = self.configured_target_base_height
|
|
|
+ self.step_count = 0
|
|
|
+ self.gait_phase = 0.0
|
|
|
+
|
|
|
+ self._rgb_renderer: mujoco.Renderer | None = None
|
|
|
+ self._viewer = None
|
|
|
+
|
|
|
+ # 提前构造一个默认起始姿态，后面 reset 时在这个基础上再叠加少量噪声。
|
|
|
+ self.initial_qpos = self.data.qpos.copy()
|
|
|
+ self.initial_qvel = np.zeros_like(self.data.qvel)
|
|
|
+ self.initial_qpos[0:3] = np.array([0.0, 0.0, self.base_height], dtype=np.float64)
|
|
|
+ self.initial_qpos[3:7] = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float64)
|
|
|
+ self.initial_qpos[self.joint_qpos_indices] = self.nominal_joint_targets.astype(np.float64)
|
|
|
+ self.data.ctrl[:] = self.nominal_joint_targets.astype(np.float64)
|
|
|
+ mujoco.mj_forward(self.model, self.data)
|
|
|
+
|
|
|
+ def _normalized_joint_position(self, joint_position: np.ndarray) -> np.ndarray:
|
|
|
+ return (joint_position - self.joint_mid) / self.joint_half_range
|
|
|
+
|
|
|
+ def _extract_snapshot(self) -> RobotStateSnapshot:
|
|
|
+ joint_position = self.data.qpos[self.joint_qpos_indices].astype(np.float32).copy()
|
|
|
+ joint_velocity = self.data.qvel[self.joint_qvel_indices].astype(np.float32).copy()
|
|
|
+ base_position = self.data.qpos[0:3].astype(np.float32).copy()
|
|
|
+ base_quat_wxyz = self.data.qpos[3:7].astype(np.float32).copy()
|
|
|
+ # MuJoCo 的自由关节四元数是 wxyz，而整个 RL 工程里统一使用 xyzw。
|
|
|
+ base_quaternion = np.array(
|
|
|
+ [base_quat_wxyz[1], base_quat_wxyz[2], base_quat_wxyz[3], base_quat_wxyz[0]],
|
|
|
+ dtype=np.float32,
|
|
|
+ )
|
|
|
+ return RobotStateSnapshot(
|
|
|
+ sim_time=float(self.data.time),
|
|
|
+ joint_position=joint_position,
|
|
|
+ joint_velocity=joint_velocity,
|
|
|
+ base_position=base_position,
|
|
|
+ base_quaternion=base_quaternion,
|
|
|
+ )
|
|
|
+
|
|
|
+ def _reference_gait_offsets(self) -> np.ndarray:
|
|
|
+ """生成一个和 Gazebo 版相同思路的参考步态。"""
|
|
|
+ offsets = np.zeros(len(self.joint_names), dtype=np.float32)
|
|
|
+ if not self.reference_gait_enabled:
|
|
|
+ return offsets
|
|
|
+
|
|
|
+ hip_amplitude = float(self.reference_gait_config.get('hip_pitch_amplitude', 0.0))
|
|
|
+ hip_bias = float(self.reference_gait_config.get('hip_pitch_bias', 0.0))
|
|
|
+ knee_amplitude = float(self.reference_gait_config.get('knee_pitch_amplitude', 0.0))
|
|
|
+ knee_bias = float(self.reference_gait_config.get('knee_pitch_bias', 0.0))
|
|
|
+ swing_knee_scale = float(self.reference_gait_config.get('swing_knee_scale', 1.0))
|
|
|
+ ankle_amplitude = float(self.reference_gait_config.get('ankle_pitch_amplitude', 0.0))
|
|
|
+ ankle_bias = float(self.reference_gait_config.get('ankle_pitch_bias', 0.0))
|
|
|
+ push_off_ankle_scale = float(self.reference_gait_config.get('push_off_ankle_scale', 0.0))
|
|
|
+ stance_ratio = float(np.clip(self.reference_gait_config.get('stance_ratio', 0.62), 0.05, 0.95))
|
|
|
+
|
|
|
+ def gait_profile(phase: float) -> tuple[float, float, float]:
|
|
|
+ phase = phase % 1.0
|
|
|
+ if phase < stance_ratio:
|
|
|
+ stance_progress = phase / stance_ratio
|
|
|
+ hip_profile = 1.0 - 2.0 * stance_progress
|
|
|
+ knee_profile = 0.18 * math.sin(math.pi * stance_progress)
|
|
|
+ push_off_progress = max((stance_progress - 0.60) / 0.40, 0.0)
|
|
|
+ ankle_profile = -0.70 * hip_profile + push_off_ankle_scale * push_off_progress
|
|
|
+ else:
|
|
|
+ swing_progress = (phase - stance_ratio) / (1.0 - stance_ratio)
|
|
|
+ hip_profile = -1.0 + 2.0 * swing_progress
|
|
|
+ knee_profile = swing_knee_scale * math.sin(math.pi * swing_progress)
|
|
|
+ ankle_profile = -0.45 * hip_profile - 0.25 * math.sin(math.pi * swing_progress)
|
|
|
+ return hip_profile, knee_profile, ankle_profile
|
|
|
+
|
|
|
+ left_hip_profile, left_knee_profile, left_ankle_profile = gait_profile(self.gait_phase)
|
|
|
+ right_hip_profile, right_knee_profile, right_ankle_profile = gait_profile(self.gait_phase + 0.5)
|
|
|
+
|
|
|
+ for index, joint_name in enumerate(self.joint_names):
|
|
|
+ if joint_name == 'left_hip_pitch_joint':
|
|
|
+ offsets[index] = hip_bias + hip_amplitude * left_hip_profile
|
|
|
+ elif joint_name == 'right_hip_pitch_joint':
|
|
|
+ offsets[index] = hip_bias + hip_amplitude * right_hip_profile
|
|
|
+ elif joint_name == 'left_knee_pitch_joint':
|
|
|
+ offsets[index] = knee_bias + knee_amplitude * left_knee_profile
|
|
|
+ elif joint_name == 'right_knee_pitch_joint':
|
|
|
+ offsets[index] = knee_bias + knee_amplitude * right_knee_profile
|
|
|
+ elif joint_name == 'left_ankle_pitch_joint':
|
|
|
+ offsets[index] = ankle_bias + ankle_amplitude * left_ankle_profile
|
|
|
+ elif joint_name == 'right_ankle_pitch_joint':
|
|
|
+ offsets[index] = ankle_bias + ankle_amplitude * right_ankle_profile
|
|
|
+
|
|
|
+ return offsets
|
|
|
+
|
|
|
+ def _advance_gait_phase(self) -> None:
|
|
|
+ if not self.reference_gait_enabled:
|
|
|
+ return
|
|
|
+ self.gait_phase = (self.gait_phase + self.control_dt / self.reference_gait_period) % 1.0
|
|
|
+
|
|
|
+ def _action_to_joint_targets(self, action: np.ndarray) -> np.ndarray:
|
|
|
+ clipped_action = np.clip(action, -1.0, 1.0)
|
|
|
+ reference_targets = self.nominal_joint_targets + self._reference_gait_offsets()
|
|
|
+ desired_residual = clipped_action * self.joint_half_range * self.action_scale
|
|
|
+
|
|
|
+ if self.action_smoothing > 0.0:
|
|
|
+ smoothed_residual = (
|
|
|
+ self.action_smoothing * self.current_action_residual
|
|
|
+ + (1.0 - self.action_smoothing) * desired_residual
|
|
|
+ )
|
|
|
+ else:
|
|
|
+ smoothed_residual = desired_residual
|
|
|
+
|
|
|
+ self.current_action_residual = smoothed_residual.astype(np.float32)
|
|
|
+ joint_targets = reference_targets + self.current_action_residual
|
|
|
+ return np.clip(joint_targets, self.joint_lower, self.joint_upper).astype(np.float32)
|
|
|
+
|
|
|
+ def _get_target_base_height(self, snapshot: RobotStateSnapshot) -> float:
|
|
|
+ if self.target_base_height is None:
|
|
|
+ self.target_base_height = float(snapshot.base_position[2])
|
|
|
+ return float(self.target_base_height)
|
|
|
+
|
|
|
+ def _build_observation(
|
|
|
+ self,
|
|
|
+ snapshot: RobotStateSnapshot,
|
|
|
+ previous_snapshot: RobotStateSnapshot,
|
|
|
+ previous_action: np.ndarray,
|
|
|
+ ) -> np.ndarray:
|
|
|
+ dt = max(snapshot.sim_time - previous_snapshot.sim_time, self.control_dt, 1e-3)
|
|
|
+ velocity = (snapshot.base_position - previous_snapshot.base_position) / dt
|
|
|
+ roll, pitch, _ = quaternion_xyzw_to_euler(snapshot.base_quaternion)
|
|
|
+ return np.concatenate(
|
|
|
+ [
|
|
|
+ self._normalized_joint_position(snapshot.joint_position),
|
|
|
+ snapshot.joint_velocity,
|
|
|
+ previous_action,
|
|
|
+ np.array(
|
|
|
+ [
|
|
|
+ snapshot.base_position[2],
|
|
|
+ roll,
|
|
|
+ pitch,
|
|
|
+ velocity[0],
|
|
|
+ self.task_config['target_forward_velocity'],
|
|
|
+ ],
|
|
|
+ dtype=np.float32,
|
|
|
+ ),
|
|
|
+ ],
|
|
|
+ dtype=np.float32,
|
|
|
+ )
|
|
|
+
|
|
|
+ def _advance_simulation(self) -> None:
|
|
|
+ for _ in range(self.frame_skip):
|
|
|
+ mujoco.mj_step(self.model, self.data)
|
|
|
+
|
|
|
+ def _terminated(self, snapshot: RobotStateSnapshot) -> bool:
|
|
|
+ if self.step_count <= self.termination_grace_steps:
|
|
|
+ return False
|
|
|
+
|
|
|
+ roll, pitch, _ = quaternion_xyzw_to_euler(snapshot.base_quaternion)
|
|
|
+ if abs(roll) > float(self.task_config['max_roll_rad']):
|
|
|
+ return True
|
|
|
+ if abs(pitch) > float(self.task_config['max_pitch_rad']):
|
|
|
+ return True
|
|
|
+ if float(snapshot.base_position[2]) < float(self.task_config['min_base_height']):
|
|
|
+ return True
|
|
|
+ return False
|
|
|
+
|
|
|
+ def _reset_qpos_with_noise(self) -> None:
|
|
|
+ self.data.qpos[:] = self.initial_qpos
|
|
|
+ self.data.qvel[:] = self.initial_qvel
|
|
|
+
|
|
|
+ if self.base_xy_noise_scale > 0.0:
|
|
|
+ self.data.qpos[0] += float(self.np_random.uniform(-self.base_xy_noise_scale, self.base_xy_noise_scale))
|
|
|
+ self.data.qpos[1] += float(self.np_random.uniform(-self.base_xy_noise_scale, self.base_xy_noise_scale))
|
|
|
+ if self.base_height_noise_scale > 0.0:
|
|
|
+ self.data.qpos[2] += float(
|
|
|
+ self.np_random.uniform(-self.base_height_noise_scale, self.base_height_noise_scale)
|
|
|
+ )
|
|
|
+ if self.joint_reset_noise_scale > 0.0:
|
|
|
+ joint_noise = self.np_random.uniform(
|
|
|
+ low=-self.joint_reset_noise_scale,
|
|
|
+ high=self.joint_reset_noise_scale,
|
|
|
+ size=len(self.joint_qpos_indices),
|
|
|
+ )
|
|
|
+ self.data.qpos[self.joint_qpos_indices] += joint_noise
|
|
|
+
|
|
|
+ self.data.ctrl[:] = self.nominal_joint_targets.astype(np.float64)
|
|
|
+ mujoco.mj_forward(self.model, self.data)
|
|
|
+
|
|
|
+ def reset(self, *, seed: int | None = None, options: dict[str, Any] | None = None):
|
|
|
+ super().reset(seed=seed)
|
|
|
+ self._reset_qpos_with_noise()
|
|
|
+ self.current_joint_target = self.nominal_joint_targets.copy()
|
|
|
+ self.current_action_residual = np.zeros(len(self.joint_names), dtype=np.float32)
|
|
|
+ self.gait_phase = 0.0
|
|
|
+
|
|
|
+ # MuJoCo 版 reset 也保留几拍站姿保持，让 position actuator 先把姿态拉回稳定区间。
|
|
|
+ for _ in range(max(self.reset_hold_steps, 0)):
|
|
|
+ self.data.ctrl[:] = self.current_joint_target.astype(np.float64)
|
|
|
+ self._advance_simulation()
|
|
|
+
|
|
|
+ snapshot = self._extract_snapshot()
|
|
|
+ self.target_base_height = self.configured_target_base_height
|
|
|
+ self._get_target_base_height(snapshot)
|
|
|
+ self.current_snapshot = snapshot
|
|
|
+ self.previous_snapshot = snapshot
|
|
|
+ self.previous_action = np.zeros(len(self.joint_names), dtype=np.float32)
|
|
|
+ self.step_count = 0
|
|
|
+
|
|
|
+ if self.render_mode == 'human':
|
|
|
+ self.render()
|
|
|
+
|
|
|
+ observation = self._build_observation(snapshot, snapshot, self.previous_action)
|
|
|
+ info = {
|
|
|
+ 'reset': True,
|
|
|
+ 'target_base_height': float(self.target_base_height),
|
|
|
+ 'target_forward_velocity': float(self.task_config['target_forward_velocity']),
|
|
|
+ }
|
|
|
+ return observation, info
|
|
|
+
|
|
|
+ def step(self, action: np.ndarray):
|
|
|
+ if self.current_snapshot is None or self.previous_snapshot is None:
|
|
|
+ raise RuntimeError('环境尚未 reset，不能直接 step。')
|
|
|
+
|
|
|
+ self.step_count += 1
|
|
|
+ action = np.asarray(action, dtype=np.float32)
|
|
|
+ joint_targets = self._action_to_joint_targets(action)
|
|
|
+ self.current_joint_target = joint_targets.copy()
|
|
|
+ self.data.ctrl[:] = joint_targets.astype(np.float64)
|
|
|
+ self._advance_simulation()
|
|
|
+ self._advance_gait_phase()
|
|
|
+
|
|
|
+ self.previous_snapshot = self.current_snapshot
|
|
|
+ self.current_snapshot = self._extract_snapshot()
|
|
|
+ terminated = self._terminated(self.current_snapshot)
|
|
|
+ truncated = self.step_count >= int(self.training_config['max_episode_steps'])
|
|
|
+
|
|
|
+ reward, reward_terms = self.reward_calculator.compute(
|
|
|
+ RewardContext(
|
|
|
+ current=self.current_snapshot,
|
|
|
+ previous=self.previous_snapshot,
|
|
|
+ action=action,
|
|
|
+ previous_action=self.previous_action,
|
|
|
+ joint_limits=self.joint_limits,
|
|
|
+ target_forward_velocity=float(self.task_config['target_forward_velocity']),
|
|
|
+ target_base_height=float(self.target_base_height),
|
|
|
+ control_dt=self.control_dt,
|
|
|
+ terminated=terminated,
|
|
|
+ )
|
|
|
+ )
|
|
|
+ observation = self._build_observation(
|
|
|
+ self.current_snapshot,
|
|
|
+ self.previous_snapshot,
|
|
|
+ self.previous_action,
|
|
|
+ )
|
|
|
+ self.previous_action = action.copy()
|
|
|
+
|
|
|
+ if self.render_mode == 'human':
|
|
|
+ self.render()
|
|
|
+
|
|
|
+ info = {
|
|
|
+ 'reward_terms': reward_terms,
|
|
|
+ 'joint_targets': joint_targets,
|
|
|
+ }
|
|
|
+ return observation, reward, terminated, truncated, info
|
|
|
+
|
|
|
+ def render(self):
|
|
|
+ if self.render_mode == 'rgb_array':
|
|
|
+ if self._rgb_renderer is None:
|
|
|
+ self._rgb_renderer = mujoco.Renderer(self.model, height=self.render_height, width=self.render_width)
|
|
|
+ self._rgb_renderer.update_scene(self.data, camera='overview')
|
|
|
+ return self._rgb_renderer.render()
|
|
|
+
|
|
|
+ if self.render_mode == 'human':
|
|
|
+ if self._viewer is None:
|
|
|
+ from mujoco import viewer
|
|
|
+
|
|
|
+ self._viewer = viewer.launch_passive(self.model, self.data)
|
|
|
+ if self._viewer is not None:
|
|
|
+ self._viewer.sync()
|
|
|
+ return None
|
|
|
+
|
|
|
+ return None
|
|
|
+
|
|
|
+ def close(self) -> None:
|
|
|
+ if self._rgb_renderer is not None:
|
|
|
+ self._rgb_renderer.close()
|
|
|
+ self._rgb_renderer = None
|
|
|
+
|
|
|
+ if self._viewer is not None:
|
|
|
+ close_method = getattr(self._viewer, 'close', None)
|
|
|
+ if callable(close_method):
|
|
|
+ close_method()
|
|
|
+ self._viewer = None
|
