190 lines
8.4 KiB
Python
190 lines
8.4 KiB
Python
import gymnasium as gym
|
||
import numpy as np
|
||
from stable_baselines3 import PPO
|
||
from stable_baselines3.common.vec_env import DummyVecEnv # 使用DummyVecEnv避免多进程问题
|
||
from stable_baselines3.common.env_util import make_vec_env
|
||
from typing import Callable, Any
|
||
from typing import Any, Literal
|
||
import numpy as np
|
||
import pybullet as p
|
||
from gymnasium import spaces
|
||
from PyFlyt.core.aviary import Aviary
|
||
from PyFlyt.core.utils.compile_helpers import check_numpy
|
||
from stable_baselines3.common.vec_env import SubprocVecEnv, DummyVecEnv
|
||
|
||
import gymnasium
|
||
import PyFlyt.gym_envs
|
||
import numpy as np
|
||
from stable_baselines3 import PPO
|
||
from stable_baselines3.common.evaluation import evaluate_policy
|
||
|
||
|
||
# --------------------------
|
||
# 核心：综合Wrapper（解决不动+调高度+自定义奖励）
|
||
# --------------------------
|
||
class QuadXPoleFullWrapper(gymnasium.Wrapper):
|
||
def __init__(
|
||
self,
|
||
env,
|
||
hover_bias=0.2, # 基础悬停PWM（解决不动：必须>0.5才够升力）
|
||
action_scale=0.2, # 动作微调范围（控制电机微调幅度，避免过大/过小）
|
||
target_height=2.5, # 目标悬停高度（调高默认高度，可改3.0/4.0）
|
||
reward_scaling=0.1 # 奖励缩放（避免奖励值过大导致训练不稳定）
|
||
):
|
||
super().__init__(env)
|
||
self.hover_bias = hover_bias # 基础悬停推力（确保无人机能起飞）
|
||
self.action_scale = action_scale# 动作微调范围（[-scale, +scale]）
|
||
self.target_height = target_height # 目标高度
|
||
self.reward_scaling = reward_scaling# 奖励缩放系数
|
||
|
||
def reset(self, **kwargs):
|
||
"""重置时将无人机初始高度设为目标高度"""
|
||
obs, info = self.env.reset(** kwargs)
|
||
# 修改无人机初始z轴位置（PyFlyt无人机状态的第3个元素是高度）
|
||
if hasattr(self.env.unwrapped, "drone"):
|
||
self.env.unwrapped.drone.state[2] = self.target_height # z轴=目标高度
|
||
return obs, info
|
||
|
||
def step(self, action):
|
||
"""1. 处理动作（确保有足够升力）；2. 自定义奖励；3. 返回新状态"""
|
||
# 1. 动作映射：模型输出[-1,1] → 实际PWM[hover_bias-scale, hover_bias+scale]
|
||
# 保证电机有基础悬停推力，解决“不动”问题
|
||
action = action * self.action_scale + self.hover_bias
|
||
# 限制动作在[0,1]（避免PWM超出物理范围导致报错）
|
||
action = np.clip(action, 0.0, 1.0)
|
||
|
||
# 2. 执行动作，获取原始环境反馈
|
||
obs, _, term, trunc, info = self.env.step(action)
|
||
|
||
# 3. 解析观测值（按PyFlyt QuadX-Pole-Balance-v3观测空间定义）
|
||
pos = obs[:3] # 无人机位置 (x, y, z)
|
||
orn = obs[3:7] # 无人机姿态（四元数 x, y, z, w）
|
||
pole_angle = obs[10] # 杆倾斜角度（核心平衡指标，索引10为主要倾斜角）
|
||
# （可选）如果需要更精准，可查看官方文档：观测空间包含杆的多个角度，取影响最大的一个
|
||
|
||
# 4. 自定义奖励计算（多维度鼓励稳定）
|
||
# ① 高度奖励：越接近目标高度，奖励越高（惩罚高度误差）
|
||
height_error = pos[2] - self.target_height
|
||
height_reward = -1.5 * (height_error ** 2) # 权重1.5，误差越小奖励越高
|
||
|
||
# ② 姿态奖励：无人机越水平，奖励越高（惩罚姿态偏移）
|
||
# 四元数x/y/z越小，姿态越接近水平（w为实部，代表水平状态）
|
||
orientation_reward = -0.8 * np.sum(orn[:3] ** 2) # 权重0.8
|
||
|
||
# ③ 杆平衡奖励：杆越竖直，奖励越高（惩罚杆倾斜）
|
||
pole_reward = -2.0 * (pole_angle ** 2) # 权重2.0，杆平衡是核心任务，权重更高
|
||
|
||
# ④ 动作平滑奖励：避免电机大幅调整（惩罚过大动作）
|
||
action_penalty = -0.1 * np.sum(action ** 2) # 权重0.1，抑制动作波动
|
||
|
||
# ⑤ 存活奖励：每步给固定奖励，鼓励持续存活（核心目标是“尽可能久”）
|
||
alive_bonus = 1.2 # 每步+1.2，存活越久总奖励越高
|
||
|
||
# 总奖励：加权求和 + 缩放
|
||
total_reward = (
|
||
height_reward + orientation_reward + pole_reward + action_penalty + alive_bonus
|
||
) * self.reward_scaling
|
||
|
||
# 5. 返回处理后的结果
|
||
return obs, total_reward, term, trunc, info
|
||
|
||
|
||
# --------------------------
|
||
# 1. 创建并包装环境
|
||
# --------------------------
|
||
# 原始环境配置（按官方文档，render_mode="human"实时显示）
|
||
|
||
# --------------------------
|
||
env_id = "PyFlyt/QuadX-Pole-Balance-v4"
|
||
|
||
# 用 make_vec_env 创建多个环境（n_envs 是并行环境数量）
|
||
env = make_vec_env(
|
||
env_id,
|
||
n_envs=4, # 4个环境同时运行（可根据CPU核心数调整）
|
||
wrapper_class=QuadXPoleFullWrapper, # 我们的自定义包装器
|
||
env_kwargs={
|
||
"render_mode": None, # 多环境训练时先不渲染，加快速度
|
||
"max_duration_seconds": 30.0,
|
||
"flight_dome_size": 5.0,
|
||
"angle_representation": "quaternion"
|
||
},
|
||
wrapper_kwargs={
|
||
"hover_bias": 0.2,
|
||
"action_scale": 0.3,
|
||
"target_height": 2.5,
|
||
"reward_scaling": 0.1
|
||
}
|
||
)
|
||
# 查看环境空间（确认配置正确）
|
||
print("动作空间（4个电机PWM）：", env.action_space)
|
||
print("观测空间（无人机+杆状态）：", env.observation_space)
|
||
|
||
|
||
# --------------------------
|
||
# 2. 定义PPO模型（适合连续动作，收敛快）
|
||
# --------------------------
|
||
model = PPO(
|
||
policy="MlpPolicy", # 多层感知器（处理连续动作）
|
||
env=env,
|
||
verbose=1, # 训练时打印详细信息（loss、reward等）
|
||
tensorboard_log="./quadx_log/", # 日志保存路径（可在TensorBoard查看训练曲线）
|
||
learning_rate=3e-4, # 学习率（连续动作任务常用3e-4）
|
||
n_steps=2048, # PPO每批收集2048步数据
|
||
batch_size=64, # 每批数据分64个batch训练（2048÷64=32，整除）
|
||
n_epochs=10, # 每批数据训练10轮
|
||
gamma=0.99, # 折扣因子（重视长期奖励）
|
||
gae_lambda=0.95, # GAE参数（平衡偏差和方差）
|
||
clip_range=0.2, # PPO裁剪范围（经典值0.2）
|
||
ent_coef=0.01, # 熵系数（鼓励探索，避免过早收敛到局部最优）
|
||
device="auto" # 自动使用GPU/CPU（有GPU会自动调用）
|
||
)
|
||
|
||
|
||
# --------------------------
|
||
# 3. 训练模型
|
||
# --------------------------
|
||
print("\n=== 开始训练 ===")
|
||
model.learn(
|
||
total_timesteps=300000, # 总训练步数（30万步，该任务较复杂，需足够步数）
|
||
log_interval=10, # 每10个批次打印一次训练信息
|
||
progress_bar=True # 显示训练进度条
|
||
)
|
||
|
||
# 保存训练好的模型（后续可直接加载，不用重新训练）
|
||
model.save("quadx_pole_balance_trained_model")
|
||
print("\n=== 模型已保存为：quadx_pole_balance_trained_model ===")
|
||
|
||
|
||
# --------------------------
|
||
# 4. 评估训练效果
|
||
# --------------------------
|
||
print("\n=== 开始评估（5局平均奖励） ===")
|
||
mean_reward, std_reward = evaluate_policy(
|
||
model=model,
|
||
env=env,
|
||
n_eval_episodes=5, # 评估5局
|
||
render=True, # 评估时实时显示
|
||
deterministic=True # 用确定性策略（避免随机动作，体现真实训练效果）
|
||
)
|
||
print(f"评估结果：平均奖励 = {mean_reward:.2f} ± {std_reward:.2f}")
|
||
# （说明：平均奖励越高、标准差越小，模型越稳定；若平均存活时间接近30秒，说明训练成功）
|
||
|
||
|
||
# --------------------------
|
||
# 5. 手动测试（可视化训练成果）
|
||
# --------------------------
|
||
print("\n=== 开始手动测试（持续1000步） ===")
|
||
obs, _ = env.reset() # 重置环境
|
||
for step in range(1000):
|
||
# 模型预测动作（确定性策略）
|
||
action, _ = model.predict(obs, deterministic=True)
|
||
# 执行动作
|
||
obs, reward, term, trunc, info = env.step(action)
|
||
# 若终止（坠毁/杆落地/超时），重置环境继续测试
|
||
if term or trunc:
|
||
print(f"第{step+1}步终止，重置环境...")
|
||
obs, _ = env.reset()
|
||
|
||
# 关闭环境（释放资源）
|
||
env.close()
|
||
print("\n=== 测试结束 ===") |