|
|
@@ -8,7 +8,7 @@ import mujoco
|
|
|
import numpy as np
|
|
|
|
|
|
from ..config import resolve_project_path
|
|
|
-from ..math_utils import quaternion_xyzw_to_euler
|
|
|
+from ..math_utils import quaternion_xyzw_to_euler, wrap_to_pi
|
|
|
from ..mujoco_model import build_mujoco_model_spec
|
|
|
from ..rewards import BipedRewardCalculator, RewardContext
|
|
|
from ..state_types import RobotStateSnapshot
|
|
|
@@ -27,6 +27,7 @@ class MujocoBipedEnv(gym.Env):
|
|
|
self.sim_config = config['sim']
|
|
|
self.task_config = config['task']
|
|
|
self.training_config = config['training']
|
|
|
+ self.commands_config = config.get('commands', {})
|
|
|
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', {})
|
|
|
@@ -52,6 +53,36 @@ class MujocoBipedEnv(gym.Env):
|
|
|
1e-4,
|
|
|
)
|
|
|
self.termination_grace_steps = int(self.task_config.get('termination_grace_steps', 0))
|
|
|
+ self.commands_enabled = bool(self.commands_config.get('enabled', False))
|
|
|
+ forward_velocity_range = self.commands_config.get(
|
|
|
+ 'forward_velocity_range',
|
|
|
+ [self.task_config['target_forward_velocity'], self.task_config['target_forward_velocity']],
|
|
|
+ )
|
|
|
+ yaw_rate_range = self.commands_config.get(
|
|
|
+ 'yaw_rate_range',
|
|
|
+ [self.task_config.get('target_yaw_rate', 0.0), self.task_config.get('target_yaw_rate', 0.0)],
|
|
|
+ )
|
|
|
+ self.command_forward_velocity_range = np.sort(np.asarray(forward_velocity_range, dtype=np.float32))
|
|
|
+ self.command_yaw_rate_range = np.sort(np.asarray(yaw_rate_range, dtype=np.float32))
|
|
|
+ self.command_resample_interval_steps = int(self.commands_config.get('resample_interval_steps', 0))
|
|
|
+ self.command_stand_probability = float(np.clip(self.commands_config.get('stand_probability', 0.0), 0.0, 1.0))
|
|
|
+ self.command_turn_only_probability = float(
|
|
|
+ np.clip(self.commands_config.get('turn_only_probability', 0.0), 0.0, 1.0)
|
|
|
+ )
|
|
|
+ self.command_combined_probability = float(
|
|
|
+ np.clip(self.commands_config.get('combined_probability', 0.0), 0.0, 1.0)
|
|
|
+ )
|
|
|
+ self.command_turn_only_forward_abs_max = float(
|
|
|
+ max(self.commands_config.get('turn_only_forward_abs_max', 0.02), 0.0)
|
|
|
+ )
|
|
|
+ self.fixed_command_config = self.commands_config.get('fixed_command', {})
|
|
|
+ self.fixed_command_enabled = bool(self.fixed_command_config.get('enabled', False))
|
|
|
+ self.command_velocity_scale = max(float(self.reference_gait_config.get('command_velocity_scale', 0.20)), 1e-6)
|
|
|
+ self.command_yaw_rate_scale = max(float(self.reference_gait_config.get('command_yaw_rate_scale', 0.30)), 1e-6)
|
|
|
+ self.command_turn_gain = float(np.clip(self.reference_gait_config.get('command_turn_gain', 0.28), 0.0, 0.8))
|
|
|
+ self.min_command_gait_scale = float(
|
|
|
+ np.clip(self.reference_gait_config.get('min_command_gait_scale', 0.0), 0.0, 1.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)))
|
|
|
@@ -90,7 +121,7 @@ class MujocoBipedEnv(gym.Env):
|
|
|
self.observation_space = gym.spaces.Box(
|
|
|
low=-np.inf,
|
|
|
high=np.inf,
|
|
|
- shape=(len(self.joint_names) * 3 + 5,),
|
|
|
+ shape=(len(self.joint_names) * 3 + (8 if self.commands_enabled else 5),),
|
|
|
dtype=np.float32,
|
|
|
)
|
|
|
|
|
|
@@ -108,6 +139,9 @@ class MujocoBipedEnv(gym.Env):
|
|
|
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.current_target_forward_velocity = float(self.task_config['target_forward_velocity'])
|
|
|
+ self.current_target_yaw_rate = float(self.task_config.get('target_yaw_rate', 0.0))
|
|
|
+ self.command_steps_until_resample = 0
|
|
|
self.step_count = 0
|
|
|
self.gait_phase = 0.0
|
|
|
|
|
|
@@ -159,6 +193,26 @@ class MujocoBipedEnv(gym.Env):
|
|
|
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))
|
|
|
+ gait_scale = 1.0
|
|
|
+ direction_sign = 1.0
|
|
|
+ left_stride_scale = 1.0
|
|
|
+ right_stride_scale = 1.0
|
|
|
+
|
|
|
+ if self.commands_enabled:
|
|
|
+ # 命令条件化训练里，让参考步态跟着命令强度逐渐放大：
|
|
|
+ # 原地命令几乎不摆腿，前进/后退/转弯命令越强，步态引导越明显。
|
|
|
+ normalized_speed = abs(self.current_target_forward_velocity) / self.command_velocity_scale
|
|
|
+ normalized_turn = abs(self.current_target_yaw_rate) / self.command_yaw_rate_scale
|
|
|
+ gait_scale = float(np.clip(max(normalized_speed, normalized_turn, self.min_command_gait_scale), 0.0, 1.0))
|
|
|
+
|
|
|
+ if self.current_target_forward_velocity < -0.03:
|
|
|
+ direction_sign = -1.0
|
|
|
+
|
|
|
+ normalized_yaw_command = float(
|
|
|
+ np.clip(self.current_target_yaw_rate / self.command_yaw_rate_scale, -1.0, 1.0)
|
|
|
+ )
|
|
|
+ left_stride_scale = max(0.2, 1.0 - self.command_turn_gain * normalized_yaw_command)
|
|
|
+ right_stride_scale = max(0.2, 1.0 + self.command_turn_gain * normalized_yaw_command)
|
|
|
|
|
|
def gait_profile(phase: float) -> tuple[float, float, float]:
|
|
|
phase = phase % 1.0
|
|
|
@@ -180,17 +234,25 @@ class MujocoBipedEnv(gym.Env):
|
|
|
|
|
|
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
|
|
|
+ offsets[index] = hip_bias + direction_sign * gait_scale * left_stride_scale * hip_amplitude * left_hip_profile
|
|
|
elif joint_name == 'right_hip_pitch_joint':
|
|
|
- offsets[index] = hip_bias + hip_amplitude * right_hip_profile
|
|
|
+ offsets[index] = (
|
|
|
+ hip_bias + direction_sign * gait_scale * right_stride_scale * hip_amplitude * right_hip_profile
|
|
|
+ )
|
|
|
elif joint_name == 'left_knee_pitch_joint':
|
|
|
- offsets[index] = knee_bias + knee_amplitude * left_knee_profile
|
|
|
+ offsets[index] = knee_bias + gait_scale * left_stride_scale * knee_amplitude * left_knee_profile
|
|
|
elif joint_name == 'right_knee_pitch_joint':
|
|
|
- offsets[index] = knee_bias + knee_amplitude * right_knee_profile
|
|
|
+ offsets[index] = knee_bias + gait_scale * right_stride_scale * knee_amplitude * right_knee_profile
|
|
|
elif joint_name == 'left_ankle_pitch_joint':
|
|
|
- offsets[index] = ankle_bias + ankle_amplitude * left_ankle_profile
|
|
|
+ offsets[index] = (
|
|
|
+ ankle_bias
|
|
|
+ + direction_sign * gait_scale * left_stride_scale * ankle_amplitude * left_ankle_profile
|
|
|
+ )
|
|
|
elif joint_name == 'right_ankle_pitch_joint':
|
|
|
- offsets[index] = ankle_bias + ankle_amplitude * right_ankle_profile
|
|
|
+ offsets[index] = (
|
|
|
+ ankle_bias
|
|
|
+ + direction_sign * gait_scale * right_stride_scale * ankle_amplitude * right_ankle_profile
|
|
|
+ )
|
|
|
|
|
|
return offsets
|
|
|
|
|
|
@@ -221,30 +283,128 @@ class MujocoBipedEnv(gym.Env):
|
|
|
self.target_base_height = float(snapshot.base_position[2])
|
|
|
return float(self.target_base_height)
|
|
|
|
|
|
+ def _base_motion_features(
|
|
|
+ self,
|
|
|
+ snapshot: RobotStateSnapshot,
|
|
|
+ previous_snapshot: RobotStateSnapshot,
|
|
|
+ ) -> tuple[float, float, float, float, float]:
|
|
|
+ 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, yaw = quaternion_xyzw_to_euler(snapshot.base_quaternion)
|
|
|
+ _, _, previous_yaw = quaternion_xyzw_to_euler(previous_snapshot.base_quaternion)
|
|
|
+ yaw_rate = wrap_to_pi(yaw - previous_yaw) / dt
|
|
|
+ return float(velocity[0]), float(velocity[1]), float(yaw_rate), float(roll), float(pitch)
|
|
|
+
|
|
|
+ def _sample_uniform_range(self, value_range: np.ndarray) -> float:
|
|
|
+ return float(self.np_random.uniform(float(value_range[0]), float(value_range[1])))
|
|
|
+
|
|
|
+ def _set_command_target(self, forward_velocity: float, yaw_rate: float) -> None:
|
|
|
+ self.current_target_forward_velocity = float(forward_velocity)
|
|
|
+ self.current_target_yaw_rate = float(yaw_rate)
|
|
|
+ self.command_steps_until_resample = int(self.command_resample_interval_steps)
|
|
|
+
|
|
|
+ def _sample_command_target(self) -> None:
|
|
|
+ if not self.commands_enabled:
|
|
|
+ self.current_target_forward_velocity = float(self.task_config['target_forward_velocity'])
|
|
|
+ self.current_target_yaw_rate = float(self.task_config.get('target_yaw_rate', 0.0))
|
|
|
+ self.command_steps_until_resample = 0
|
|
|
+ return
|
|
|
+
|
|
|
+ if self.fixed_command_enabled:
|
|
|
+ self.current_target_forward_velocity = float(
|
|
|
+ self.fixed_command_config.get('forward_velocity', self.task_config['target_forward_velocity'])
|
|
|
+ )
|
|
|
+ self.current_target_yaw_rate = float(
|
|
|
+ self.fixed_command_config.get('yaw_rate', self.task_config.get('target_yaw_rate', 0.0))
|
|
|
+ )
|
|
|
+ self.command_steps_until_resample = 0
|
|
|
+ return
|
|
|
+
|
|
|
+ stand_probability = self.command_stand_probability
|
|
|
+ turn_only_probability = self.command_turn_only_probability
|
|
|
+ combined_probability = self.command_combined_probability
|
|
|
+ translate_only_probability = max(1.0 - stand_probability - turn_only_probability - combined_probability, 0.0)
|
|
|
+ total_probability = stand_probability + turn_only_probability + combined_probability + translate_only_probability
|
|
|
+ if total_probability <= 0.0:
|
|
|
+ stand_probability = 0.0
|
|
|
+ turn_only_probability = 0.0
|
|
|
+ combined_probability = 0.0
|
|
|
+ translate_only_probability = 1.0
|
|
|
+ total_probability = 1.0
|
|
|
+
|
|
|
+ sample = float(self.np_random.uniform(0.0, 1.0))
|
|
|
+ stand_cutoff = stand_probability / total_probability
|
|
|
+ turn_only_cutoff = stand_cutoff + turn_only_probability / total_probability
|
|
|
+ translate_only_cutoff = turn_only_cutoff + translate_only_probability / total_probability
|
|
|
+
|
|
|
+ if sample < stand_cutoff:
|
|
|
+ self._set_command_target(0.0, 0.0)
|
|
|
+ return
|
|
|
+ if sample < turn_only_cutoff:
|
|
|
+ forward_velocity = float(
|
|
|
+ self.np_random.uniform(-self.command_turn_only_forward_abs_max, self.command_turn_only_forward_abs_max)
|
|
|
+ )
|
|
|
+ yaw_rate = self._sample_uniform_range(self.command_yaw_rate_range)
|
|
|
+ self._set_command_target(forward_velocity, yaw_rate)
|
|
|
+ return
|
|
|
+ if sample < translate_only_cutoff:
|
|
|
+ self._set_command_target(self._sample_uniform_range(self.command_forward_velocity_range), 0.0)
|
|
|
+ return
|
|
|
+
|
|
|
+ self._set_command_target(
|
|
|
+ self._sample_uniform_range(self.command_forward_velocity_range),
|
|
|
+ self._sample_uniform_range(self.command_yaw_rate_range),
|
|
|
+ )
|
|
|
+
|
|
|
+ def _advance_command_schedule(self) -> None:
|
|
|
+ if not self.commands_enabled or self.fixed_command_enabled or self.command_resample_interval_steps <= 0:
|
|
|
+ return
|
|
|
+
|
|
|
+ self.command_steps_until_resample -= 1
|
|
|
+ if self.command_steps_until_resample <= 0:
|
|
|
+ self._sample_command_target()
|
|
|
+
|
|
|
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)
|
|
|
+ forward_velocity, lateral_velocity, yaw_rate, roll, pitch = self._base_motion_features(
|
|
|
+ snapshot,
|
|
|
+ previous_snapshot,
|
|
|
+ )
|
|
|
+ if self.commands_enabled:
|
|
|
+ base_features = np.array(
|
|
|
+ [
|
|
|
+ snapshot.base_position[2],
|
|
|
+ roll,
|
|
|
+ pitch,
|
|
|
+ forward_velocity,
|
|
|
+ lateral_velocity,
|
|
|
+ yaw_rate,
|
|
|
+ self.current_target_forward_velocity,
|
|
|
+ self.current_target_yaw_rate,
|
|
|
+ ],
|
|
|
+ dtype=np.float32,
|
|
|
+ )
|
|
|
+ else:
|
|
|
+ base_features = np.array(
|
|
|
+ [
|
|
|
+ snapshot.base_position[2],
|
|
|
+ roll,
|
|
|
+ pitch,
|
|
|
+ forward_velocity,
|
|
|
+ self.current_target_forward_velocity,
|
|
|
+ ],
|
|
|
+ dtype=np.float32,
|
|
|
+ )
|
|
|
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,
|
|
|
- ),
|
|
|
+ base_features,
|
|
|
],
|
|
|
dtype=np.float32,
|
|
|
)
|
|
|
@@ -307,6 +467,7 @@ class MujocoBipedEnv(gym.Env):
|
|
|
self.previous_snapshot = snapshot
|
|
|
self.previous_action = np.zeros(len(self.joint_names), dtype=np.float32)
|
|
|
self.step_count = 0
|
|
|
+ self._sample_command_target()
|
|
|
|
|
|
if self.render_mode == 'human':
|
|
|
self.render()
|
|
|
@@ -315,7 +476,8 @@ class MujocoBipedEnv(gym.Env):
|
|
|
info = {
|
|
|
'reset': True,
|
|
|
'target_base_height': float(self.target_base_height),
|
|
|
- 'target_forward_velocity': float(self.task_config['target_forward_velocity']),
|
|
|
+ 'target_forward_velocity': float(self.current_target_forward_velocity),
|
|
|
+ 'target_yaw_rate': float(self.current_target_yaw_rate),
|
|
|
}
|
|
|
return observation, info
|
|
|
|
|
|
@@ -335,6 +497,8 @@ class MujocoBipedEnv(gym.Env):
|
|
|
self.current_snapshot = self._extract_snapshot()
|
|
|
terminated = self._terminated(self.current_snapshot)
|
|
|
truncated = self.step_count >= int(self.training_config['max_episode_steps'])
|
|
|
+ target_forward_velocity = float(self.current_target_forward_velocity)
|
|
|
+ target_yaw_rate = float(self.current_target_yaw_rate)
|
|
|
|
|
|
reward, reward_terms = self.reward_calculator.compute(
|
|
|
RewardContext(
|
|
|
@@ -343,12 +507,14 @@ class MujocoBipedEnv(gym.Env):
|
|
|
action=action,
|
|
|
previous_action=self.previous_action,
|
|
|
joint_limits=self.joint_limits,
|
|
|
- target_forward_velocity=float(self.task_config['target_forward_velocity']),
|
|
|
+ target_forward_velocity=target_forward_velocity,
|
|
|
+ target_yaw_rate=target_yaw_rate,
|
|
|
target_base_height=float(self.target_base_height),
|
|
|
control_dt=self.control_dt,
|
|
|
terminated=terminated,
|
|
|
)
|
|
|
)
|
|
|
+ self._advance_command_schedule()
|
|
|
observation = self._build_observation(
|
|
|
self.current_snapshot,
|
|
|
self.previous_snapshot,
|
|
|
@@ -362,6 +528,8 @@ class MujocoBipedEnv(gym.Env):
|
|
|
info = {
|
|
|
'reward_terms': reward_terms,
|
|
|
'joint_targets': joint_targets,
|
|
|
+ 'target_forward_velocity': target_forward_velocity,
|
|
|
+ 'target_yaw_rate': target_yaw_rate,
|
|
|
}
|
|
|
return observation, reward, terminated, truncated, info
|
|
|
|
