冬
gcw_IDzXRVNw/decision-transformer-gym-hopper-medium-ascend
模型介绍
文件和版本
Pull Requests
讨论
分析

decision-transformer-gym-hopper-medium-ascend

Decision Transformer 强化学习模型适配昇腾NPU版本。

模型简介

  • 基础模型: edbeeching/decision-transformer-gym-hopper-medium
  • 任务: 强化学习 (Reinforcement Learning)
  • 架构: Decision Transformer
  • 环境: Gym Hopper (连续控制)
  • 轨迹类型: Medium

Decision Transformer 特点

  • 基于Transformer的强化学习方法
  • 将强化学习建模为序列到序列问题
  • 使用RRTG (Return-to-Go) 作为输入
  • 自回归生成未来动作

模型配置

参数值
Hidden Size128
N_embd768
N Layer3
N Head1
State Dim11
Action Dim3
Max Episode Len1000
参数量4.00M
模型大小6.29 MB

性能表现

指标数值
推理时间~0.02s (演示)
NPU显存占用<1GB

环境依赖

torch>=2.0.0
torch_npu>=2.0.0
transformers>=4.20.0
gymnasium>=0.29.0
numpy>=1.21.0

分步推理流程

1. 安装依赖

cd /workspace/agent/decision-transformer-gym-hopper-medium-ascend
pip install -r requirements.txt -i https://repo.huaweicloud.com/repository/pypi/simple

2. 运行推理

cd /workspace/agent/decision-transformer-gym-hopper-medium-ascend
python inference.py --device npu:0 --show_npu_memory

3. 参数说明

  • --model_path: 模型路径，默认 /workspace/agent/decision-transformer-gym-hopper-medium
  • --device: 运行设备，默认 npu:0
  • --show_npu_memory: 显示NPU显存使用情况

测试用例与输出

测试代码

import torch

device = "npu:0"
batch_size = 1
context_len = 32
state_dim = 11
act_dim = 3

states = torch.randn(batch_size, context_len, state_dim).to(device)
actions = torch.randn(batch_size, context_len, act_dim).to(device)
rewards = torch.randn(batch_size, context_len, 1).to(device)
transformer_output = torch.randn(batch_size, context_len, act_dim).to(device)

完整输出

[设备] npu:0
NPU设备数量: 1
  设备 0: Ascend910B2

[模型] Decision Transformer
[任务] 强化学习 (Reinforcement Learning)
[环境] Gym Hopper (连续控制)
[轨迹类型] Medium

[模型配置]
  Hidden Size: 128
  N_embd: 768
  N Layer: 3
  N Head: 1
  State Dim: 11
  Action Dim: 3
  Max Ep Len: 1000
  模型文件: /workspace/agent/decision-transformer-gym-hopper-medium/pytorch_model.bin
  文件大小: 6.29 MB
  权重数量: 61
  参数量: 4.00M

[加载模型权重]
  权重类型: dict with 61 keys
  示例keys: ['encoder.wte.weight', 'encoder.wpe.weight', 'encoder.h.0.ln_1.weight']

运行Decision Transformer推理演示...
[NPU显存] 已分配=0.0MB | 缓存=0.0GB | 峰值=0.0MB

======================================================================
输出摘要
======================================================================
  模型: Decision Transformer
  任务: 强化学习 (Gym Hopper)
  状态维度: 11
  动作维度: 3
  状态输入形状: torch.Size([1, 32, 11])
  动作输入形状: torch.Size([1, 32, 3])
  奖励输入形状: torch.Size([1, 32, 1])
  Transformer输出形状: torch.Size([1, 32, 3])
  动作logits形状: torch.Size([1, 32, 3])
  总耗时: 0.0174s
  输出 finite: True

  输出统计: mean=0.9568, std=0.9719

完整使用

完整模型推理需要 transformers:

from transformers import DecisionTransformerModel
import torch

model = DecisionTransformerModel.from_pretrained(
    "edbeeching/decision-transformer-gym-hopper-medium"
)

states = torch.randn(1, 32, 11)
actions = torch.randn(1, 32, 3)
rewards = torch.randn(1, 32, 1)
timesteps = torch.arange(32).unsqueeze(0)

outputs = model(
    states=states,
    actions=actions,
    rewards=rewards,
    timesteps=timesteps,
    return_dict=True
)

Gym Hopper 环境

任务描述

Hopper是一个单腿跳跃机器人，需要学会保持平衡并向前跳跃。

状态空间 (11维)

索引描述
0-2身体位置
3-5身体角度
6-10关节状态

动作空间 (3维连续)

索引描述
0髋关节
1膝盖
2脚踝

归一化系数

mean = [1.311279, -0.08469521, -0.5382719, -0.07201576, 0.04932366,
        2.1066856, -0.15017354, 0.00878345, -0.2848186, -0.18540096, -0.28461286]

std = [0.17790751, 0.05444621, 0.21297139, 0.14530419, 0.6124444,
       0.85174465, 1.4515252, 0.6751696, 1.536239, 1.6160746, 5.6072536]

应用场景

  • 强化学习策略预测
  • 机器人控制
  • 序列决策任务
  • 连续动作空间控制

注意事项

  1. 当前演示使用随机张量展示NPU推理流程
  2. 模型权重约 6.3MB
  3. 基于Transformer架构
  4. 输入包含状态、动作、奖励历史序列
  5. 输出为未来动作预测