SKRL 介绍
1. 简介
skrl 是模块化的强化学习库,支持 PyTorch、JAX、Gymnasium 和 Isaac Lab 等。其基本组件包括:
- Environment:产生观测、奖励和终止信号。
- Model:定义 PolicyNet 和 ValueNet。
- Memory:保存采集到的交互数据。
- Agent:实现 PPO 等强化学习算法。
- Trainer:管理环境交互和模型训练。
1.1. PolicyNet、ValueNet 与 Reward
- PolicyNet:输入观测,输出动作或动作分布,负责决定“怎么做”。
- ValueNet:输入观测,预测当前状态的累计回报,负责评价“当前状态有多好”。
- Reward:环境在动作执行后给出的评分,用于计算累计回报和优势值。
PPO 根据 Reward 和 ValueNet 的预测计算优势值与价值目标:
优势值 ≈ 实际累计回报 - ValueNet 预测值- 优势值为正:提高 PolicyNet 选择该动作的概率。
- 优势值为负:降低 PolicyNet 选择该动作的概率。
- PolicyNet 利用优势值更新策略。
- ValueNet 拟合由 Reward 和 GAE 计算出的价值目标,从而减小值函数误差。
1.2. Policy 输出与 Isaac Lab 执行器
Isaac Lab 中的控制流程为:
Policy 输出动作
↓
ActionManager 调用由 ActionsCfg 配置的 ActionTerm,完成缩放、偏移、限幅和关节映射
↓
Articulation 接收关节控制目标
└── 内部显式 Actuator 或 PhysX drive 计算或约束关节力矩
↓
PhysX 执行物理仿真- ActionsCfg:配置 ActionTerm 如何缩放、偏移、限幅及映射 Policy 输出,再将处理结果设置为关节位置、速度或力矩目标。
- ArticulationCfg:配置机器人模型、关节和 Actuator,不直接接收运行时动作。
- Actuator:Articulation 的组成部分;显式 Actuator 计算或约束力矩,隐式 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_1 和 joint_2。ArticulationCfg.actuators 中通过 joint_names_expr 匹配这些关节的 ActuatorCfg,负责配置相应关节的执行器,最后由 PhysX 执行。Policy 输出将经过 scale、offset 和限幅,因此不一定等于实际关节位置。
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 每步理论最大约为 1,reward / steps 越接近 1,说明杆越直、小车越靠近中心。