四足机器人控制任务

任务概述

四足机器人控制示例

HalfCheetah 是传统强化学习中的一个经典控制问题。关于该任务的详细描述请参考 Mujoco-HalfCheetah。D4RL 是一个经典的离线强化学习数据集，本示例将采用其中的 halfcheetah-medium-v2 数据集进行 REVIVE 训练。下面将介绍如何通过 REVIVE 在 halfcheetah-medium-v2 数据集上训练得到理想的环境与策略。最后，我们将对比数据集中的原始策略与 REVIVE 得到的策略，直观展示 REVIVE 的强大能力。

项目	描述
Action Space	Continuous(6,)
Observation	Shape (17,)
Observation	High [inf inf inf inf inf inf inf inf inf inf inf inf inf inf inf inf inf]
Observation	Low [-inf -inf -inf -inf -inf -inf -inf -inf -inf -inf -inf -inf -inf -inf -inf -inf -inf]

动作空间

动作空间由连续的 6 维向量组成，分别表示控制各关节的力矩。每个维度的取值范围为 [-1, 1]。

观察空间

观察空间是一个 17 维向量，包含各关节的角度、角速度，以及机器人在 X 轴、Z 轴的线速度，以及机器人的高度。

初始状态

机器人从原点开始，从初始状态准备向前奔跑。

任务结束

机器人跑完 1000 步后任务会被强制结束。

任务目标

机器人的目标是在固定步数内尽可能多地向前奔跑，同时保持自身能量消耗较小。这一目标在奖励函数的定义中得到体现。

训练流程

REVIVE 是一个基于历史数据的离线强化学习工具。在四足机器人控制任务上使用 REVIVE 的完整流程如下：

收集历史决策数据：收集四足机器人控制任务的历史运行数据
构建决策流图和训练数据：
- 结合业务场景和历史数据构建决策流图
- 决策流图使用 .yaml 文件描述业务数据的交互逻辑
- 训练数据使用 .npz 或 .h5 文件存储决策流图中定义的节点数据
定义奖励函数：
- 根据任务目标设计奖励函数
- 奖励函数指导控制策略优化，确保机器人能够高效向前奔跑
开始模型训练：
- 完成决策流图、训练数据和奖励函数定义后
- 使用 REVIVE 进行虚拟环境模型训练和策略模型训练
上线测试：
- 将训练好的策略模型部署到实际环境进行测试

准备数据

本示例无需手动收集历史数据，因为 D4RL 库已提供标准的离线历史数据。首先需要下载并预处理 D4RL 数据集，使其符合 REVIVE 的输入格式。数据处理脚本位于 data/generate_data.py，进入 data 目录运行以下命令即可得到处理后的数据集：

bash

python generate_data.py

数据处理过程中需要注意以下关键点：

轨迹切分：如准备数据所述，REVIVE 数据集需要包含 index 信息。我们通过判断数据集中 t+1 时刻的 obs 是否与 t 时刻的 next_obs 一致来切分轨迹，并生成相应的 index 信息。
还原 delta_x 信息：halfcheetah-medium-v2 数据集不直接提供 x 坐标信息，而 x 坐标信息对计算 reward 至关重要。因此，我们通过数据集中的 reward 信息来还原 delta_x 信息：

d e l t a_x := x_{t + 1} - x_{t}

具体处理细节请参考 data/generate_data.py 文件。

处理完成后得到 .npz 文件，将其放入 data/ 文件夹中。

定义决策流图

下面的示例展示了 .yaml 文件的详细配置。.yaml 文件通常包含两个主要部分：graph 和 columns。其中 graph 部分定义决策流图，columns 部分定义数据组成。具体配置方法请参考准备数据文档。

由于 obs 包含 17 个维度，obs 的列需要 按顺序 定义在 columns 部分。如 Mujoco-HalfCheetah 所示，状态和动作中的变量都是连续的，因此使用 continuous 类型描述每一列数据。

另外，注意这里对 delta_x 定义了 min、max 范围为 [-1, 1]。这是因为在策略训练时，每一步的 delta_x 可能会超出数据集的范围（数据集中 delta_x 约在 [-0.12, 0.43]）。这对 REVIVE 的 数据归一化 处理有重要影响。默认情况下，REVIVE 会读取数据中的 min、max 值进行归一化。

yaml

metadata:
  columns:
  - obs_0:
      dim: obs
      type: continuous
  - obs_1:
      dim: obs
      type: continuous
  ...
  - obs_16:
      dim: obs
      type: continuous

  - action_0:
      dim: action
      type: continuous
  - action_1:
      dim: action
      type: continuous
  ...
  - action_5:
      dim: action
      type: continuous

  - delta_x:
      dim: delta_x
      type: continuous
      min: -1
      max: 1

  graph:
    action:
    - obs
    delta_x:
    - obs
    - action
    next_obs:
    - obs
    - action
    - delta_x

这样我们就得到了 .yaml 文件，将其放入 data/ 文件夹中。

构建奖励函数

这里我们使用 Mujoco 中对 HalfCheetah 定义的奖励函数，详细实现请参考 HalfCheetah-Env

python

import torch
import numpy as np
from typing import Dict


def get_reward(data : Dict[str, torch.Tensor]) -> torch.Tensor:
    action = data["action"]
    delta_x = data["delta_x"]

    forward_reward_weight = 1.0
    ctrl_cost_weight = 0.1
    dt = 0.05

    if isinstance(action, np.ndarray):
        array_type = np
        ctrl_cost = ctrl_cost_weight * array_type.sum(array_type.square(action), axis=-1, keepdims=True)
    else:
        array_type = torch
        # ctrl_cost 代表执行动作的体能开销，由动作的二范数平方构成
        ctrl_cost = ctrl_cost_weight * array_type.sum(array_type.square(action), axis=-1, keepdim=True)

    x_velocity = delta_x / dt
    # forward_reward 代表 halfcheetah 向前运动的奖励，x_velocity 越大，奖励值越高
    forward_reward = forward_reward_weight * x_velocity

    # 最终 halfcheetah 得到的 reward 由 forward_reward 和 ctrl_cost 构成
    # 这对应了 halfcheetah 任务的目标：在固定步数内尽可能多地向前奔跑，同时保持自身能量消耗较小
    reward = forward_reward - ctrl_cost

    return reward

这样我们就得到了奖励函数文件，将其放入 data/ 文件夹中。

使用 REVIVE SDK 训练控制策略

现在我们已经构建完成运行 REVIVE SDK 所需的所有文件，包括 .npz 数据文件、.yaml 文件和 reward.py 奖励函数。另外还有 config.json 文件，保存了训练所需的超参数。这四个文件都位于 data/ 文件夹中。当前文件目录结构如下：

|-- data
|   |-- config.json
|   |-- generate_data.py
|   |-- halfcheetah_medium-v2.hdf5
|   |-- halfcheetah-medium-v2.npz
|   |-- halfcheetah-medium-v2.yaml
|   `-- halfcheetah_reward.py
`-- train.py

用户可以切换到 examples/task/HalfCheetah 目录下，运行以下 Python 命令开始虚拟环境模型训练和策略模型训练。在训练过程中，可以随时使用 TensorBoard 打开日志目录来监控训练过程。

bash

python train.py \
    -df data/halfcheetah-medium-v2.npz \
    -cf data/halfcheetah-medium-v2.yaml \
    -rf data/halfcheetah_reward.py \
    -rcf data/config.json \
    --target_policy_name action \
    -vm once \
    -pm once \
    --run_id halfcheetah-medium-v2-revive \
    --revive_epoch 1500 \
    --sac_epoch 1500

重要提示

REVIVE SDK 已经提供了训练所需的数据和代码，详情请参考 REVIVE SDK 源码库。

测试模型

当 REVIVE 完成虚拟环境模型训练和策略模型训练后，我们可以在日志文件夹（logs/<run_id>）下找到保存的模型（.pkl 或 .onnx）。我们尝试在真实环境上测试策略的效果，并与数据中的控制效果进行对比。在下面的测试代码中，我们将策略在真实环境中运行 100 轮，每轮执行 1000 步，输出这 100 次的总平均回报（累计奖励）。REVIVE SDK 的策略获得了 7156.0 平均奖励，远高于数据中策略的 4770.3 奖励值，控制效果提高了约 50%。

python

import pickle
import d4rl
import gym
import numpy as np

def take_revive_action(state):
    new_data = {}
    new_data['obs'] = state
    action = policy_revive.infer(new_data)
    return action

policy_revive = pickle.load(open('policy.pkl', 'rb'))
env = gym.make('halfcheetah-medium-v2')

re_list = []
for traj in range(100):
    state = env.reset()

    obs = state
    re_turn = []
    done = False
    while not done:
        action = take_revive_action(obs)
        next_state, reward, done, _ = env.step(action)
        obs = next_state
        re_turn.append(reward)

    print(np.sum(np.array(re_turn)[:]))
    re_list.append(np.sum(re_turn))

print('mean return:', np.mean(re_list), ' std:', np.std(re_list),
     ' normal_score:', env.get_normalized_score(np.mean(re_list)))

# REVIVE 平均回报:
# mean return: 7155.900144836804  std: 63.78200350280033  normal_score: 0.5989506173038248

为了更直观地比较策略效果，我们生成策略的控制对比动画。可以发现，REVIVE 的策略能够控制 HalfCheetah 跑得更快更稳定，比数据中的原始策略更加优秀。

四足机器人控制效果对比

四足机器人控制任务 ​

任务概述 ​

动作空间 ​

观察空间 ​

初始状态 ​

任务结束 ​

任务目标 ​

训练流程 ​

准备数据 ​

定义决策流图 ​

构建奖励函数 ​

使用 REVIVE SDK 训练控制策略 ​

测试模型 ​