AmurBear/uav-rwd

UAV Transformer Reward Model

模型描述

基于Transformer架构的无人机奖励模型，用于强化学习环境中的奖励预测。模型使用因果掩码（Causal Mask）确保step-wise推理，支持变长序列输入。

模型详情

• 模型类型: Transformer-based Reward Model
• 任务: 无人机轨迹奖励预测
• 架构: 4层Transformer，8个注意力头
• 参数量: 868,865
• 输入维度: 26维特征向量
• 训练数据: 10,677个轨迹文件，约190万数据点

训练信息

数据来源

• 仿真环境: PX4 + IsaacSim
• 数据规模:
  • 总样本数: 909,280
  • 训练集: 818,352
  • 验证集: 90,928

训练配置

• 设备: CUDA (GPU)
• 批次大小: 32
• 训练轮数: 50 epochs
• 学习率: 0.0001 (余弦退火)
• 优化器: Adam

性能指标

• 训练损失: 0.0007
• 验证损失: 0.0000
• MSE: 0.0998
• MAE: 0.2904
• 相关系数: 0.898

输入特征 (26维)

特征	维度	说明
相对位置	3	相对首帧的位置 (x, y, z)
相对目标位置	3	相对首帧的目标位置
位置误差	3	当前位置与目标位置的差
线速度	3	无人机线速度
目标速度	3	目标速度
线加速度	3	无人机线加速度
姿态	3	滚转、俯仰、偏航角
角速度	3	角速度
相对时间	1	相对首帧的时间
位置误差范数	1	位置误差的L2范数

使用方法

安装依赖

pip install torch numpy

加载模型

from transformer_reward_model import TransformerRewardModel
import numpy as np

# 加载模型
model = TransformerRewardModel()
model.load_model('transformer_reward_model_full.pth')

# 准备输入 (30个时间步，26维特征)
sequence = np.random.randn(30, 26)

# 预测奖励
rewards = model.predict(sequence, use_causal_mask=True)
current_reward = rewards[-1, 0]

RL框架集成

from rl_inference import RewardModelInference

# 初始化推理接口
reward_model = RewardModelInference(
    model_path='transformer_reward_model_full.pth',
    history_length=30,
    batch_size=16
)

# 单个预测
reward = reward_model.predict_single(features)

# 批次预测
rewards = reward_model.predict_batch(batch_features)

VERL框架集成

from demos.verl_custom_reward_manager import create_drone_reward_manager

# 创建VERL兼容的奖励管理器
reward_manager = create_drone_reward_manager(
    model_path='transformer_reward_model_full.pth',
    history_length=30,
    batch_size=32
)

核心特性

1. 因果掩码: 确保预测时只能看到当前及之前的信息
2. Step-wise推理: 支持任意长度的历史状态序列
3. 相对首帧位置: 所有位置特征相对首帧归一化
4. 变长序列支持: 支持10-500步的输入序列
5. 批次推理: 支持批次大小1-32+

限制

• 模型在PX4+IsaacSim仿真环境训练，真实环境可能需要微调
• 输入必须是26维特征向量
• 建议历史窗口长度为30步

引用

如果使用本模型，请引用：

@misc{uav_transformer_reward_model_2026,
  title={UAV Transformer Reward Model},
  author={Your Name},
  year={2026},
  howpublished={\url{https://huggingface.co/your-username/uav-reward-model}}
}

许可证

MIT License

联系方式

• GitHub: [your-github-username]
• Email: [your-email]