SKRL 介绍

1. 简介

skrl 是模块化的强化学习库,支持 PyTorch、JAX、Gymnasium 和 Isaac Lab 等。其基本组件包括:

1.1. PolicyNet、ValueNet 与 Reward

PPO 根据 Reward 和 ValueNet 的预测计算优势值与价值目标:

优势值 ≈ 实际累计回报 - ValueNet 预测值

1.2. Policy 输出与 Isaac Lab 执行器

Isaac Lab 中的控制流程为:

Policy 输出动作
    ↓
ActionManager 调用由 ActionsCfg 配置的 ActionTerm,完成缩放、偏移、限幅和关节映射
    ↓
Articulation 接收关节控制目标
    └── 内部显式 Actuator 或 PhysX drive 计算或约束关节力矩
    ↓
PhysX 执行物理仿真

例如 JointEffortActionCfg(scale=100.0) 收到 Policy 输出 0.5 时,在不考虑 offset 和限幅的情况下,生成名义值为 50 的关节力命令,再由 Actuator 和 PhysX 处理。

1.2.1. 多关节动作的对应关系

Policy 不直接选择 Actuator,而是通过关节名称建立两级映射:

Policy 动作向量索引
    ↓ ActionsCfg.joint_names
Articulation 中的关节
    ↓ ActuatorCfg.joint_names_expr
管理该关节的 Actuator

例如:

JointPositionActionCfg(
    asset_name="robot",
    joint_names=["joint_1", "joint_2"],
    preserve_order=True,
)

此时 Policy 输出 [0.2, -0.4] 分别对应 joint_1joint_2ArticulationCfg.actuators 中通过 joint_names_expr 匹配这些关节的 ActuatorCfg,负责配置相应关节的执行器,最后由 PhysX 执行。Policy 输出将经过 scaleoffset 和限幅,因此不一定等于实际关节位置。


2. 训练

2.1. 安装依赖

pip install torch tensorboard \
  "skrl==2.1.0" \
  "gymnasium[classic-control]==1.2.1"

2.2. skrl_train.py

from pathlib import Path
from typing import cast

import gymnasium as gym
import torch.nn as nn
from skrl.agents.torch.ppo import PPO, PPO_CFG
from skrl.envs.wrappers.torch import wrap_env
from skrl.memories.torch import RandomMemory
from skrl.models.torch import CategoricalMixin, DeterministicMixin, Model
from skrl.trainers.torch import SequentialTrainer


POLICY_PATH = Path(__file__).with_name("checkpoints") / "cartpole_policy.pt"


# 自定义 Reward:
# 1. 杆越直,奖励越高;
# 2. 小车越靠近中心,奖励越高;
# 3. 倒下时额外惩罚。
class CustomRewardWrapper(gym.Wrapper):
    def step(self, action):
        observation, _, terminated, truncated, info = self.env.step(action)

        position, _, angle, angular_velocity = observation

        reward = (
            1.0
            - 2.0 * abs(angle) / 0.2095
            - 0.1 * abs(position) / 2.4
            - 0.01 * abs(angular_velocity)
        )

        if terminated:
            reward -= 5.0

        info["custom_reward"] = reward
        return observation, float(reward), terminated, truncated, info


class Policy(CategoricalMixin, Model):
    def __init__(self, observation_space, action_space, device):
        Model.__init__(
            self,
            observation_space=observation_space,
            action_space=action_space,
            device=device,
        )
        CategoricalMixin.__init__(self, unnormalized_log_prob=True)

        self.network = nn.Sequential(
            # 全连接输入层:把观测向量映射为 64 个隐藏特征。
            # CartPole 的 self.num_observations=4,因此该层结构为 4 -> 64。
            nn.Linear(self.num_observations, 64),
            nn.Tanh(),
            nn.Linear(64, 64),
            nn.Tanh(),
            nn.Linear(64, self.num_actions),
        )

    def compute(self, inputs, role):
        # 输出两个离散动作的 logits
        return self.network(inputs["observations"]), {}


class Value(DeterministicMixin, Model):
    def __init__(self, observation_space, action_space, device):
        Model.__init__(
            self,
            observation_space=observation_space,
            action_space=action_space,
            device=device,
        )
        DeterministicMixin.__init__(self)

        self.network = nn.Sequential(
            nn.Linear(self.num_observations, 64),
            nn.Tanh(),
            nn.Linear(64, 64),
            nn.Tanh(),
            nn.Linear(64, 1),
        )

    def compute(self, inputs, role):
        return self.network(inputs["observations"]), {}


def main() -> None:
    # 创建并包装 Gymnasium 环境
    env = wrap_env(CustomRewardWrapper(gym.make("CartPole-v1")))
    observation_space = cast(gym.Space, env.observation_space)
    action_space = cast(gym.Space, env.action_space)

    models = {
        "policy": Policy(
            # CartPole 示例:Box(shape=(4,)),依次表示位置、速度、杆角度和杆角速度。
            observation_space=observation_space,
            # CartPole 示例:Discrete(2),动作 0 表示向左推,动作 1 表示向右推。
            action_space=action_space,
            device=env.device,
        ).to(env.device),
        "value": Value(
            observation_space, action_space, env.device
        ).to(env.device),
    }

    cfg = PPO_CFG(
        # 每次更新前,每个并行环境收集 1024 步交互数据。
        rollouts=1024,
        learning_epochs=8,
        mini_batches=8,
        learning_rate=3e-4,
        entropy_loss_scale=0.01,
    )

    memory = RandomMemory(
        memory_size=cfg.rollouts,
        num_envs=env.num_envs,
        device=env.device,
    )

    agent = PPO(
        models=models,
        memory=memory,
        cfg=cfg,
        observation_space=observation_space,
        action_space=action_space,
        device=env.device,
    )

    trainer = SequentialTrainer(
        cfg={"timesteps": 50_000, "headless": True},
        env=env,
        agents=agent,
    )

    try:
        trainer.train()

        # 部署只需要 PolicyNet,因此训练完成后将其单独保存到固定路径。
        POLICY_PATH.parent.mkdir(parents=True, exist_ok=True)
        models["policy"].save(str(POLICY_PATH))
        print(f"PolicyNet 已保存到:{POLICY_PATH}")
    finally:
        env.close()


if __name__ == "__main__":
    main()

3. 推理

3.1. skrl_play.py

import argparse
from typing import cast

import gymnasium as gym
import torch
from skrl.envs.wrappers.torch import wrap_env
from skrl_train import CustomRewardWrapper, POLICY_PATH, Policy


def main() -> None:
    parser = argparse.ArgumentParser(
        description="使用训练好的 PolicyNet 运行 CartPole"
    )
    parser.add_argument("--episodes", type=int, default=5, help="推理回合数")
    parser.add_argument(
        "--render", action="store_true",
        help="显示 Gymnasium 动画窗口"
    )
    args = parser.parse_args()

    if not POLICY_PATH.exists():
        raise FileNotFoundError(
            f"找不到 PolicyNet:{POLICY_PATH},请先运行 python skrl_train.py"
        )

    render_mode = "human" if args.render else None
    env = wrap_env(
        CustomRewardWrapper(gym.make("CartPole-v1", render_mode=render_mode))
    )
    observation_space = cast(gym.Space, env.observation_space)
    action_space = cast(gym.Space, env.action_space)

    # 推理只构造 PolicyNet,不需要 ValueNet、Memory 或 PPO 训练器。
    policy = Policy(
        observation_space=observation_space,
        action_space=action_space,
        device=env.device,
    ).to(env.device)
    policy.load(str(POLICY_PATH))
    policy.eval()

    try:
        for episode in range(1, args.episodes + 1):
            observations, _ = env.reset()
            episode_reward = 0.0
            episode_steps = 0

            while True:
                with torch.inference_mode():
                    # PolicyNet 输出两个动作的 logits;推理时选择 logits 最大的动作。
                    logits, _ = policy.compute(
                        {"observations": observations}, role="policy"
                    )
                    actions = torch.argmax(logits, dim=-1, keepdim=True)

                observations, rewards, terminated, truncated, _ = env.step(
                    actions
                )
                episode_reward += float(rewards.item())
                episode_steps += 1

                if bool((terminated | truncated).any().item()):
                    break

            print(
                f"Episode {episode}: steps={episode_steps}, "
                f"reward={episode_reward:.2f}"
            )
    finally:
        env.close()


if __name__ == "__main__":
    main()

当前自定义 Reward 每步理论最大约为 1reward / steps 越接近 1,说明杆越直、小车越靠近中心。