Eval 与 Export¶

GeneLab 的研究复现工具链 —— genelab.rl.evaluator / eval_callback / exporter，对应三条 CLI，把 train → eval → export 闭环串起来：

命令	用途	输出
`genelab eval TASK CKPT`	Deterministic rollout，固定 seed，N episode	`eval.json`
`genelab train ... --eval-every K`	训练期周期 eval + 保存 best model	`logs/.../best_model.<ext>` + `best_model_meta.json`
`genelab export TASK CKPT`	不依赖后端的 TorchScript / ONNX policy	`policy.{ts,onnx}` + `<file>.metadata.json`

三个命令都走同一个 backend 抽象（InferenceSetup，定义在 genelab.rl.backends.base），所以在 rsl_rl / skrl / sb3 上行为一致。

`genelab eval`¶

跑 vectorized deterministic rollout，并按下面的 schema 写一份 JSON：

genelab eval GeneLab-Inverted-Pendulum-v0 logs/rsl_rl/exp1/.../model_500.pt \
    --num-envs 64 --episodes 100 --seed 0 \
    --deterministic --out eval.json

输出：

{
  "task": "GeneLab-Inverted-Pendulum-v0",
  "checkpoint": "logs/.../model_500.pt",
  "num_episodes": 100,
  "metrics": {
    "return_mean": 487.3,
    "return_std": 22.1,
    "length_mean": 998.4,
    "success_rate": 0.96
  },
  "wall_clock_seconds": 18.2,
  "seed": 0,
  "deterministic": true,
  "evaluated_at": "2026-05-20T08:42:11+00:00"
}

成功率¶

success_rate 在任务通过 ManagerBasedRlEnv.step 的 extras["is_success"] publish 一个 per-env bool tensor 时计算（gymnasium 约定）。任务在 termination / reward 项中设置 self._extras["is_success"] = <(num_envs,) bool tensor> — 通常是 manipulation 的目标姿态判定、locomotion 的速度命令跟踪判定等。

任务没暴露 is_success 时输出 success_rate: null；下游工具（best-model 选择、reference-runs 表）必须容忍 None。

`genelab train --eval-every`¶

设了 --eval-every K 后，训练按 K 个 iteration 分 chunk 跑。每 chunk 结束后，加载最新 checkpoint 进同一 backend，跑一次 deterministic eval（默认 10 episode，num_envs 与训练相同）。当 return_mean 超过历史最佳，checkpoint 被复制到 <log_dir>/best_model.<ext>，同时 best_model_meta.json 写入 eval payload。

genelab train GeneLab-Inverted-Pendulum-v0 \
    --max_iterations 1000 --num_envs 64 --seed 0 \
    --eval-every 100 --eval-episodes 16

注意：

每个 chunk 走 backend 的正常 train lifecycle，会关闭并重建 Genesis env。短任务把 --eval-every 设到 ≥ 50，让 Genesis 初始化时间被摊薄。
Off-policy 算法（skrl / sb3 的 SAC / TD3 / DDPG）每 chunk 重载 checkpoint 会丢 replay buffer，sample efficiency 下降但仍能收敛。
best_model.<ext> 复用来源 backend 的 checkpoint 格式（rsl_rl / skrl 用 .pt，sb3 用 .zip）。meta 文件记录来源 iter、eval seed、episode 数、return 统计。

`genelab export`¶

把 actor 子网络序列化为 TorchScript 或 ONNX，把每个 obs 项的 scale / clip 烘焙进单一 forward(raw_obs) -> actions 调用。部署侧只需 torch（TorchScript）或一个 ONNX runtime，不需要 rsl_rl / skrl / stable_baselines3。

# TorchScript
genelab export Genelab-Velocity-Flat-Unitree-G1-v0 logs/.../model_30000.pt \
    --format torchscript --out policy.ts

# ONNX（默认 opset 17）
genelab export Genelab-Velocity-Flat-Unitree-G1-v0 logs/.../model_30000.pt \
    --format onnx --out policy.onnx --opset 17

注：GeneLab-Franka-Pick-And-Place-v0 是 SAC+HER + goal-conditioned Dict 观测。它导出的模型输入是单一扁平 obs，即 observation + achieved_goal + desired_goal 按该顺序拼接（见下方多 group 的 metadata）。 Locomotion 任务（cartpole / G1）则是单一 flat-tensor obs group。

导出器会在旁边写一份 <output>.metadata.json，描述 obs schema。obs_dim 是扁平输入的总宽度；每个 obs_groups 条目记录它在扁平张量中的 start 偏移（便于把 goal-conditioned policy 切回各子空间）：

{
  "task": "Genelab-Velocity-Flat-Unitree-G1-v0",
  "checkpoint": "logs/.../model_30000.pt",
  "obs_dim": 23,
  "obs_groups": {
    "policy": {
      "start": 0,
      "dim": 23,
      "terms": [
        {"name": "joint_pos", "dim": 7, "start": 0, "scale": 1.0, "clip": null},
        {"name": "joint_vel", "dim": 7, "start": 7, "scale": 0.1, "clip": [-2, 2]}
      ]
    }
  },
  "action_dim": 7,
  "action_range": [-1.0, 1.0],
  "normalization_baked": true,
  "is_recurrent": false,
  "format": "torchscript",
  "exported_at": "2026-05-20T08:42:11+00:00",
  "torch_version": "2.4.0"
}

SAC+HER 任务的 obs_groups 会有三个条目 —— 例如 observation（start: 0）、 achieved_goal（start: 35）、desired_goal（start: 38）—— obs_dim 为它们之和。

部署侧用法¶

import torch
m = torch.jit.load("policy.ts")
m.eval()
# raw obs（按训练时的拼接顺序）；模型自己应用 scale/clip
actions = m(torch.tensor([[joint_pos_0, joint_pos_1, ..., joint_vel_0, ...]]))

ONNX：

import onnxruntime as ort
sess = ort.InferenceSession("policy.onnx")
actions = sess.run(None, {"obs": raw_obs.astype("float32")})[0]

实际导出的内容¶

actor 通过 backend 各自的小 shim 取出来，包成统一的调用形态：

rsl_rl：直接从算法对象取 actor 模块（alg._raw_actor，没有则退回 alg.actor），并用它的 as_jit() 导出包装——后者暴露扁平的 forward(obs) -> 确定性动作，且已把学到的 obs normalizer 烘焙进去。仍兼容把 actor 放在 alg.actor_critic.actor（或只有 act_inference）的旧版本。
skrl：包 agent.policy.act，对 GaussianMixin policy 返回 deterministic mean（mean_actions key）。
sb3：包 model.policy._predict(obs, deterministic=True)，对 PPO / A2C / SAC / TD3 / DDPG 统一。对目标条件化的 SAC+HER policy，observation 是 Dict（observation / achieved_goal / desired_goal），SAC actor 会用到所有 key，因此 wrapper 接收这些子空间按该顺序拼接的扁平张量，在调用 policy 前重建出 Dict —— 导出的模型仍是单一扁平 obs 输入，metadata 的 obs_groups 记录每个子空间的 start / dim，部署方据此还原布局。

SB3 的 ONNX 导出使用旧版基于 TorchScript 的导出器（torch.onnx.export(..., dynamo=False)）：基于 torch.export 的新默认导出器（torch ≥ 2.9）无法 trace SAC 的 Normal 分布构造。

Recurrent（RNN / LSTM / GRU）policy 导出¶

在 RslRlModelCfg 上设置 rnn_type（"lstm" 或 "gru"）即训练一个循环策略 —— 这是唯一开关，会自动选用 rsl_rl 的 RNNModel：

RslRlModelCfg(rnn_type="lstm", rnn_hidden_dim=256, rnn_num_layers=1)

genelab export 随后自动走循环导出路径（无需额外参数）。metadata 会多出 "is_recurrent": true 以及一个 "recurrent" 块，记录 rnn_type、rnn_num_layers、 rnn_hidden_dim、hidden state 形状和 ONNX 端口名。

两种格式对 hidden state 的暴露方式不同：

TorchScript 把 hidden state 藏在模块内部，因此调用形态仍是单输入的 MLP 形式 forward(obs) -> actions。模块还额外暴露 reset() 方法 —— 每个 episode 边界都要调用它来清零 hidden state。序列化的 buffer 固定 batch size 为 1（单台部署机器人）。

import torch
m = torch.jit.load("policy.ts"); m.eval()
m.reset()                       # 每个 episode 开始时
actions = m(raw_obs)            # 喂 raw obs；hidden state 内部维护

ONNX 显式暴露 hidden state。输入是 obs、h_in（LSTM 还有 c_in）；输出是 actions、h_out（LSTM 还有 c_out），形状均为 (num_layers, batch, hidden_dim)。每步把返回的状态回灌，并在 episode 边界清零：

import numpy as np, onnxruntime as ort
sess = ort.InferenceSession("policy.onnx")
h = np.zeros((num_layers, 1, hidden_dim), np.float32)
c = np.zeros((num_layers, 1, hidden_dim), np.float32)   # 仅 LSTM
actions, h, c = sess.run(None, {"obs": raw_obs, "h_in": h, "c_in": c})
# GRU: actions, h = sess.run(None, {"obs": raw_obs, "h_in": h})

play / eval rollout 会对刚结束 episode 的环境自动清零 hidden state，因此循环策略的 eval 指标不会被跨 episode 的残留状态污染。

限制¶

导出的模型不应用 ObservationTermCfg.noise —— noise 只在训练时启用。