企业官网建设流程全解析

Stable Baselines3 模型导出实战：从 PyTorch 策略到 ONNX、C++、TensorFlow.js 与 Coral 的多框架部署指南

【免费下载链接】stable-baselines3PyTorch version of Stable Baselines, reliable implementations of reinforcement learning algorithms.项目地址: https://gitcode.com/GitHub_Trending/st/stable-baselines3

训练完成一个强化学习智能体之后，真正的挑战往往才开始：如何把它部署到另一个语言、另一个推理框架，甚至是浏览器、移动端或边缘设备上？本指南以 Stable Baselines3（SB3）官方导出文档为核心，完整讲解策略（Policy）作为控制器的导出原理，并给出从 ONNX、PyTorch JIT（C++）、ONNX Runtime Web、TensorFlow.js 到 TFLite/Coral（Edge TPU）的全链路可运行代码，以及基于state_dict/get_parameters的手动导出和 SBX（SB3 + JAX）经 PyTorch 中转导出的完整示例。读完本文，你将能够在脱离 Gym 与 SB3 运行时的前提下，在任何目标框架中完成智能体的推理。

背景：策略（Policy）就是控制器

在 Stable Baselines3 中，真正决定"看到什么状态、输出什么动作"的控制器存放在**策略（policy）**对象内部。每个学习算法（DQN、A2C、PPO、SAC、TD3、DDPG 等）都持有一个 policy 对象，代表当前学到的行为，可通过model.policy访问。策略中保存了推理（predict action）所需的全部信息——网络结构、权重参数、观察预处理逻辑——因此"导出策略"就等价于"导出控制器"。

这一点在源码中非常清晰：以 policy 基类 为例，ActorCriticPolicy.forward()（stable_baselines3/common/policies.py#L636）对 PPO 这类 actor-critic 算法返回(actions, values, log_prob)三元组，即动作、状态值和动作对数概率；而所有推理的"糖衣"逻辑（观察转换、图像归一化、动作后处理）都在predict()方法中（stable_baselines3/common/policies.py#L331-L386）。导出时我们绕开predict()，直接对 policy 模块本身做 trace 或转换为中间格式，就能在其他框架复现同样的前向计算。

提示：可结合 examples.md 中的训练示例理解策略在完整训练流程中的角色。

导出前必须弄清的两处隐式处理

直接导出 policy 时，有两处"藏在内部"的处理如果不手动补齐，导出的模型在目标框架中会产生错误输出。

CNN 观察的归一化（内置于 policy 内部）

当使用 CNN 策略处理图像观察时，观察会在预处理阶段被归一化：除以 255，把像素值缩放到 [0, 1]。这个预处理发生在 policy 内部（preprocess_obs，见 stable_baselines3/common/preprocessing.py#L91-L139），因此导出后的模型会自动包含该归一化逻辑，外部无需再除一次 255。这也意味着：如果你在导出前手动把图像除以 255，再送入导出的模型，就会得到双重归一化。建议用原始 0–255 的像素值作为模型输入。

# 预处理核心逻辑（摘录自 stable_baselines3/common/preprocessing.py） if normalize_images and is_image_space(observation_space): return obs.float() / 255.0

另外要注意：PyTorch 使用channel-first（通道在前）的布局，Gymnasium 的图像观察通常是HWC（高度、宽度、通道）。若需要从外部传入图像，可能需要先转置为 channel-first（maybe_transpose在 stable_baselines3/common/policies.py#L236-L277 的obs_to_tensor中处理了这一逻辑）。

连续动作的后处理（clip 或 unscale）

第二个隐藏步骤是连续动作的后处理。predict()在得到网络原始输出后，还会对动作做一步修正（stable_baselines3/common/policies.py#L372-L379）：

• 若策略启用了squash_output（SAC 的 actor 默认用 tanh 把动作压到 [-1, 1]），则调用unscale_action把动作从 [-1, 1] 重新映射回action_space的实际范围（如 Pendulum-v1 的 [-2, 2]）；
• 否则直接把动作np.clip到action_space.low/action_space.high，避免高斯采样越界。

这两个后处理步骤都不会出现在你直接导出的 ONNX/JIT 模型里，导出的模型返回的是"归一化/未缩放的"原始网络输出。因此，在目标框架推理后，必须由你自己补上 clip 或 unscale 逻辑（下文各小节会给出对应代码）。

导出到 ONNX（PyTorch 2.0+ / Opset 14+）

如果使用 PyTorch 2.0+ 与 ONNX Opset 14+，可以非常轻松地把 SB3 策略导出为 ONNX。核心思路是：把model.policy包进一个th.nn.Module包装类，在forward中直接调用 policy 并固定deterministic=True（若要导出随机策略，则改为deterministic=False）。

再次强调：以下导出的模型不包含连续动作的后处理步骤（clip 或 unscale 到正确的动作空间）。

PPO 完整导出示例

以 PPO 为例（Pendulum-v1，MlpPolicy），完整代码如下——训练、保存、加载、包装、导出、用 onnxruntime 校验一气呵成：

import torch as th from typing import Tuple from stable_baselines3 import PPO from stable_baselines3.common.policies import BasePolicy class OnnxableSB3Policy(th.nn.Module): def __init__(self, policy: BasePolicy): super().__init__() self.policy = policy def forward(self, observation: th.Tensor) -> Tuple[th.Tensor, th.Tensor, th.Tensor]: # NOTE: Preprocessing is included, but postprocessing # (clipping/inscaling actions) is not, # If needed, you also need to transpose the images so that they are channel first # use deterministic=False if you want to export the stochastic policy # policy() returns `actions, values, log_prob` for PPO return self.policy(observation, deterministic=True) # Example: model = PPO("MlpPolicy", "Pendulum-v1") PPO("MlpPolicy", "Pendulum-v1").save("PathToTrainedModel") model = PPO.load("PathToTrainedModel.zip", device="cpu") onnx_policy = OnnxableSB3Policy(model.policy) observation_size = model.observation_space.shape dummy_input = th.randn(1, *observation_size) th.onnx.export( onnx_policy, dummy_input, "my_ppo_model.onnx", opset_version=17, input_names=["input"], ) ##### Load and test with onnx import onnx import onnxruntime as ort import numpy as np onnx_path = "my_ppo_model.onnx" onnx_model = onnx.load(onnx_path) onnx.checker.check_model(onnx_model) observation = np.zeros((1, *observation_size)).astype(np.float32) ort_sess = ort.InferenceSession(onnx_path) actions, values, log_prob = ort_sess.run(None, {"input": observation}) print(actions, values, log_prob) # Check that the predictions are the same with th.no_grad(): print(model.policy(th.as_tensor(observation), deterministic=True))

几个值得注意的细节：

• dummy input 的 batch 维度：th.randn(1, *observation_size)显式带上 batch 维1，这是 ONNX 导出所必需的固定形状。导出后如果用 batch size 为 1 推理，可直接使用；需要动态 batch 时可另行配置dynamic_axes。
• 加载到 CPU：PPO.load(..., device="cpu")避免 GPU 环境导出后在其他机器上因缺少 CUDA 而无法加载。
• 校验一致性：用onnx.checker.check_model做结构合法性检查，再用相同输入对比 ONNX Runtime 与原始 PyTorch policy 的输出，确保转换无误。由于 PPO 是 actor-critic 结构，forward返回三个输出（actions、values、log_prob），ONNX 模型的输出也是三个。

对于MultiInputPolicy（字典观察，如 goal-conditioned 任务），导出方式类似，社区在相关 issue 中有更详细的讨论（可参考 GH#1873 中的方案）。

SAC：只导出 actor 网络

对 SAC，过程类似，但示例中只导出 actor 网络——因为 rollout 时 actor 足以决定动作，不需要 critic。SAC 的 actor 输出是 tanh 压到 [-1, 1] 的缩放动作，因此后处理（unscale）必不可少：

import torch as th from stable_baselines3 import SAC class OnnxablePolicy(th.nn.Module): def __init__(self, actor: th.nn.Module): super().__init__() self.actor = actor def forward(self, observation: th.Tensor) -> th.Tensor: # NOTE: You may have to postprocess (unnormalize) actions # to the correct bounds (see commented code below) return self.actor(observation, deterministic=True) # Example: model = SAC("MlpPolicy", "Pendulum-v1") SAC("MlpPolicy", "Pendulum-v1").save("PathToTrainedModel.zip") model = SAC.load("PathToTrainedModel.zip", device="cpu") onnxable_model = OnnxablePolicy(model.policy.actor) observation_size = model.observation_space.shape dummy_input = th.randn(1, *observation_size) th.onnx.export( onnxable_model, dummy_input, "my_sac_actor.onnx", opset_version=17, input_names=["input"], ) ##### Load and test with onnx import onnxruntime as ort import numpy as np onnx_path = "my_sac_actor.onnx" observation = np.zeros((1, *observation_size)).astype(np.float32) ort_sess = ort.InferenceSession(onnx_path) scaled_action = ort_sess.run(None, {"input": observation})[0] print(scaled_action) # Post-process: rescale to correct space # Rescale the action from [-1, 1] to [low, high] # low, high = model.action_space.low, model.action_space.high # post_processed_action = low + (0.5 * (scaled_action + 1.0) * (high - low)) # Check that the predictions are the same with th.no_grad(): print(model.actor(th.as_tensor(observation), deterministic=True))

这里model.policy.actor就是 SAC 策略中的 actor 模块（见 stable_baselines3/sac/policies.py 中SACPolicy通过self.make_actor()构建 actor 的实现，stable_baselines3/sac/policies.py#L281）。onnxable_model的输入是观察、输出是 [-1, 1] 区间的缩放动作；拿到scaled_action后，需按如下公式还原到动作空间真实范围：

low, high = model.action_space.low, model.action_space.high post_processed_action = low + (0.5 * (scaled_action + 1.0) * (high - low))

这一公式与 SB3 内部的unscale_action完全一致（stable_baselines3/common/policies.py#L402-L413）。

导出到 C++（PyTorch JIT Trace）

如果目标是把模型嵌入 C++ 推理代码，可以用 PyTorch JIT 对模型进行trace、freeze 并优化，保存为 TorchScript 文件，然后在任意支持 TorchScript 的环境（包括 libtorch C++）中加载推理：

# See "ONNX export" for imports and OnnxablePolicy jit_path = "sac_traced.pt" # Trace and optimize the module traced_module = th.jit.trace(onnxable_model.eval(), dummy_input) frozen_module = th.jit.freeze(traced_module) frozen_module = th.jit.optimize_for_inference(frozen_module) th.jit.save(frozen_module, jit_path) ##### Load and test with torch import torch as th dummy_input = th.randn(1, *observation_size) loaded_module = th.jit.load(jit_path) action_jit = loaded_module(dummy_input)

要点说明：

• th.jit.trace(module.eval(), dummy_input)：基于一组 dummy 输入记录执行路径生成计算图。务必先.eval()，把 BatchNorm/Dropout 切换到推理模式，保证 trace 出的图稳定。
• th.jit.freeze：将权重固化进图中、移除梯度信息，进一步减小体积。
• th.jit.optimize_for_inference：针对推理做算子融合等优化，适合部署场景。
• th.jit.load后再喂入同样的dummy_input，即可在纯 TorchScript 环境下得到与 ONNX 一致的动作输出（同样不含后处理）。

社区（RL Zoo 项目）中已有 C++ 导出相关的草案实现可供参考，思路与本小节完全一致：trace → freeze → 在 C++ 侧加载.pt。

导出到 ONNX-JS / ONNX Runtime Web（浏览器推理）

把模型跑在浏览器里是常见的部署需求，官方推荐的做法是：先用前面的方法导出 ONNX，再用onnxruntime-web在浏览器中加载推理。社区有完整的端到端示例（一个"躲避汽车"的驾驶环境），流程为：

1. 创建/训练一个 PPO 模型；
2. 将模型导出为 ONNX，同时把归一化统计量（normalization stats）存成 JSON；
3. 在浏览器中用onnxruntime-web加载 ONNX 模型，并在前端做同样的归一化，达到与本地训练近似的效果。

下面的最小示例演示"转 ONNX → 浏览器无后处理推理"（以 SAC 为例，导出部分与上文相同）：

import torch as th from stable_baselines3 import SAC class OnnxablePolicy(th.nn.Module): def __init__(self, actor: th.nn.Module): super().__init__() self.actor = actor def forward(self, observation: th.Tensor) -> th.Tensor: # NOTE: You may have to postprocess (unnormalize or renormalize) return self.actor(observation, deterministic=True) # Example: model = SAC("MlpPolicy", "Pendulum-v1") SAC("MlpPolicy", "Pendulum-v1").save("PathToTrainedModel.zip") model = SAC.load("PathToTrainedModel.zip", device="cpu") onnxable_model = OnnxablePolicy(model.policy.actor) observation_size = model.observation_space.shape dummy_input = th.randn(1, *observation_size) th.onnx.export( onnxable_model, dummy_input, "my_sac_actor.onnx", opset_version=17, input_names=["input"], )

前端 JavaScript 推理（依赖通过npm install onnxruntime-web安装，测试版本为 1.19，也可用 CDN 引入）：

// Install using `npm install onnxruntime-web` (tested with version 1.19) or using cdn import * as ort from 'onnxruntime-web'; async function runInference() { const session = await ort.InferenceSession.create('my_sac_actor.onnx'); // The observation_size = 3 (for Pendulum-v1) const inputData = Float32Array.from([0.1, -0.2, 0.3]); const inputTensor = new ort.Tensor('float32', inputData, [1, 3]); const results = await session.run({ input: inputTensor }); const outputName = session.outputNames[0]; const action = results[outputName].data; console.log('Predicted action=', action); } runInference();

注意这里input名称必须与th.onnx.export时指定的input_names=["input"]一致；Pendulum-v1 的观察维度为 3，因此输入张量形状是[1, 3]。同样的，输出动作是缩放/未归一化的原始值，业务侧如需落到真实动作空间仍需自行后处理。

导出到 TensorFlow.js

把模型跑在 TensorFlow.js 中需要走一条较长的转换链：SB3（Torch）⇒ ONNX ⇒ TensorFlow ⇒ TensorFlow.js，中间涉及多个工具链的版本兼容问题，因此本节给出经过验证的完整方案。

注意（截至 2025 年 11 月的信息）：onnx2tf尚不支持 TensorFlow.js，因此必须改用tfjs-converter。但tfjs-converter目前维护不活跃，要求使用较旧的 opset 与 TensorFlow 版本。

关键约束是ONNX 的 opset 版本必须改为 14（上文 ONNX 导出示例中opset_version=17是为了更高版本下的更稳定用法，而此处需要opset_version=14）。

第一步仍是 SB3 ⇒ ONNX（与上文一致，把opset_version改为 14）。随后在全新环境中安装指定版本的依赖（已在 Python 3.10 下测试通过）：

pip install --use-deprecated=legacy-resolver tensorflow==2.13.0 keras==2.13.1 onnx==1.16.0 onnx-tf==1.9.0 tensorflow-probability==0.21.0 tensorflowjs==4.15.0 jax==0.4.26 jaxlib==0.4.26

然后执行 ONNX ⇒ TensorFlow 的转换：

import onnx import onnx_tf.backend import tensorflow as tf ONNX_FILE_PATH = "my_sac_actor.onnx" MODEL_PATH = "tf_model" onnx_model = onnx.load(ONNX_FILE_PATH) onnx.checker.check_model(onnx_model) print(onnx.helper.printable_graph(onnx_model.graph)) print('Converting ONNX to TF...') tf_rep = onnx_tf.backend.prepare(onnx_model) tf_rep.export_graph(MODEL_PATH) # After this do not forget to use `tensorflowjs_converter`

若目录结构正确且无报错，再执行命令行转换（把 TensorFlow SavedModel 转成 tfjs_graph_model）：

tensorflowjs_converter --input_format=tf_saved_model --output_format=tfjs_graph_model tf_model tfjs_model

如果tensorflowjs_converter报错，先升级 TensorFlow 相关包（pip install --upgrade tensorflow tensorflow-decision-forests tensorflowjs）再重试，通常即可成功（无需重跑上一步转换代码）。

前端加载tfjs_model目录中的model.json进行推理：

import * as tf from 'https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@4.15.0/+esm'; // Post processing not included async function runInference() { const MODEL_URL = './tfjs_model/model.json'; const model = await tf.loadGraphModel(MODEL_URL); // Observation_size is 3 for Pendulum-v1 const inputData = [1.0, 0.0, 0.0]; const inputTensor = tf.tensor2d([inputData], [1, 3]); const resultTensor = model.execute(inputTensor); const action = await resultTensor.data(); console.log('Predicted action=', action); inputTensor.dispose(); resultTensor.dispose(); } runInference();

tf.loadGraphModel加载的是图模型（graph model），用model.execute(inputTensor)执行前向并读取结果。注意代码同样省略了动作后处理，且观察输入[1, 3]对应 Pendulum-v1 的 3 维状态。

导出到 TFLite / Coral（Edge TPU）

Google 为边缘端 AI 部署推出了Coral芯片，它有多种形态（包括 USB 版本）。把 SB3 训练的模型跑在树莓派 + Coral USB 上，正是社区相关示例的最初动机。Coral 芯片推理快、功耗极低，但设备端训练能力有限，所以需要先把网络量化到与 Coral 能力匹配的形式。

从 SB3 到 Coral 的完整链路为：

SB3 (Torch) ⇒ ONNX ⇒ TensorFlow ⇒ TFLite ⇒ Coral

社区有完整的最小示例，覆盖了整条导出链路并演示了多数导出变体的前向推理，同时专门处理了以下容易踩坑的问题：

• 让 Gym 的观察能够正确适配 ONNX（形状、dtype、归一化）；
• 合理量化 TFLite 模型，既对齐 Gym 的动作语义，又充分利用 Coral 的加速能力；
• 复用前文介绍的OnnxablePolicy包装类完成第一步导出。

整体思路与 TensorFlow.js 一节类似：先在opset_version=14下完成 SB3 ⇒ ONNX，再用onnx-tf转成 TensorFlow，最后用 TFLite 转换器生成量化模型并部署到 Coral 设备。

手动导出（state_dict 与 get_parameters）

如果不依赖任何转换工具，也可以手动导出需要的参数（权重），在目标框架中自行重建网络。SB3 提供了两种取参方式：

1. model.get_parameters()：返回 agent 全部网络的 state-dict 映射（按对象名组织的字典，见 stable_baselines3/common/base_class.py#L804-L817）。它基于_get_torch_save_params()收集所有需要保存的模块（策略、critic 等）并逐一导出state_dict()，如果还需要访问优化器的状态字典，就必须用这个接口。
2. model.policy.state_dict()：policy 本身也是 PyTorchnn.Module，可以直接调用标准的state_dict()拿到网络参数。

关于架构信息：每个算法的网络结构请查看各自目录下的policies.py，例如 stable_baselines3/ppo/policies.py、stable_baselines3/sac/policies.py、stable_baselines3/dqn/policies.py 等，里面定义了 MLP 层数、激活函数、特征提取器（如NatureCNN、MlpExtractor）等结构细节，据此即可在目标框架中用相同层序重建前向计算。

建议：大多数情况下优先使用 PyTorch 标准的state_dict()/load_state_dict()（只涉及网络参数），只有当确实需要优化器状态时才改用get_parameters()。

SBX（SB3 + JAX）导出到 ONNX

作为手动导出的典型示例，Stable Baselines Jax（SBX）的策略可以借助"中间 PyTorch 表示"导出到 ONNX：先把 JAX 的 Flax 参数字典映射成 PyTorchstate_dict，重建等价网络，再走标准th.onnx.export。这样不仅演示了跨框架取参，也示范了 SBX 中 actor 网络的内部结构。

import numpy as np import sbx import torch as th class TorchPolicy(th.nn.Module): def __init__(self, obs_dim: int, hidden_dim: int, act_dim: int): super().__init__() self.net = th.nn.Sequential( th.nn.Linear(obs_dim, hidden_dim), th.nn.Tanh(), th.nn.Linear(hidden_dim, hidden_dim), th.nn.Tanh(), th.nn.Linear(hidden_dim, act_dim), ) def forward(self, x: th.Tensor) -> th.Tensor: return self.net(x) model = sbx.PPO("MlpPolicy", "Pendulum-v1") # Also possible: load a trained model # model = sbx.PPO.load("PathToTrainedModel.zip") params = model.policy.actor_state.params["params"] # For debug: print("=== SBX params ===") for key, value in params.items(): if isinstance(value, dict): for name, val in value.items(): print(f"{key}.{name}: {val.shape}", end=" ") else: print(f"{key}: {value.shape}", end=" ") print("\n" + "=" * 20 + "\n") obs_dim = model.observation_space.shape act_dim = model.action_space.shape # Number of units in the hidden layers (assume a network architecture like [64, 64]) hidden_dim = params["Dense_0"]["kernel"].shape[1] # map params to torch state_dict keys num_layers = len([k for k in params.keys() if k.startswith("Dense_")]) state_dict = {} for i in range(num_layers): layer_name = f"Dense_{i}" state_dict[f"net.{i * 2}.bias"] = th.from_numpy(np.array(params[layer_name]["bias"])) state_dict[f"net.{i * 2}.weight"] = th.from_numpy(np.array(params[layer_name]["kernel"].T)) torch_policy = TorchPolicy(obs_dim[0], hidden_dim, act_dim[0]) print("=== Torch params ===") print(" ".join(f"{key}:{tuple(value.shape)}" for key, value in torch_policy.named_parameters())) print("=" * 20 + "\n") torch_policy.load_state_dict(state_dict) torch_policy.eval() dummy_input = th.zeros((1, *obs_dim)) # Use normal Torch export th.onnx.export( torch_policy, (dummy_input,), "my_ppo_actor.onnx", opset_version=18, input_names=["input"], output_names=["action"], ) ##### Load and test with onnx import onnxruntime as ort onnx_path = "my_ppo_actor.onnx" ort_sess = ort.InferenceSession(onnx_path) observation = np.random.random((1, *obs_dim)).astype(np.float32) action = ort_sess.run(None, {"input": observation})[0] print(action) sbx_action, _ = model.predict(observation, deterministic=True) with th.no_grad(): torch_action = torch_policy(th.as_tensor(observation)) # Check that the predictions are the same assert np.allclose(sbx_action, action) assert np.allclose(sbx_action, torch_action.numpy())

这段代码有三个关键点值得展开：

• 参数来源：SBX 的 actor 权重位于model.policy.actor_state.params["params"]，按 Flax 惯例以Dense_0、Dense_1等命名层，每个 Dense 层包含kernel与bias两个数组。打印各参数形状便于对照调试。
• 映射细节：Flax 的kernel形状是(in_features, out_features)，而 PyTorchnn.Linear的weight形状是(out_features, in_features)，因此必须做转置（kernel.T）；bias 直接复制。重建的网络结构（两层 64 隐层 + Tanh 激活 + 输出层）要严格对应 SBX PPO 默认的[64, 64]架构，hidden_dim从params["Dense_0"]["kernel"].shape[1]动态读取。
• 一致性校验：load_state_dict后分别跑 SBX 原生predict、PyTorch 重建网络、ONNX Runtime 三个前向，用np.allclose断言三者输出一致，确保"跨框架取参 → 重建 → 转 ONNX"整条链路没有偏差。

小结

Stable Baselines3 的模型导出本质上围绕一个事实展开：策略对象持有推理所需的全部信息。从导出目标反推，选择对应的技术路径即可：

无论走哪条路径，请牢记两点：CNN 策略的图像归一化（除以 255）已在 policy 内部完成，外部不要再重复归一化；连续动作的 clip/unscale 后处理不会随模型导出，需要在目标框架的推理侧自行实现。掌握这些原则后，SB3 训练出的智能体即可无缝迁移到 C++、浏览器、移动端与边缘设备，真正打通"训练—部署"的最后一步。

【免费下载链接】stable-baselines3PyTorch version of Stable Baselines, reliable implementations of reinforcement learning algorithms.项目地址: https://gitcode.com/GitHub_Trending/st/stable-baselines3